camel-ai 0.2.67__py3-none-any.whl → 0.2.80a2__py3-none-any.whl

camel/toolkits/vertex_ai_veo_toolkit.py

@@ -0,0 +1,590 @@
+# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
+
+import asyncio
+import base64
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+from camel.logger import get_logger
+from camel.toolkits import FunctionTool
+from camel.toolkits.base import BaseToolkit
+from camel.utils import api_keys_required, dependencies_required
+
+logger = get_logger(__name__)
+
+
+@api_keys_required(
+    [("GOOGLE_APPLICATION_CREDENTIALS", "GOOGLE_APPLICATION_CREDENTIALS")]
+)
+class VertexAIVeoToolkit(BaseToolkit):
+    r"""A toolkit for interacting with Google Vertex AI Veo video generation.
+
+    This toolkit provides methods for generating videos using Google's Veo,
+    supporting both text-to-video and image-to-video generation with various
+    customization options.
+    """
+
+    @dependencies_required('google.cloud.aiplatform')
+    def __init__(
+        self,
+        project_id: Optional[str] = None,
+        location: str = "us-central1",
+        model_id: str = "veo-2.0-generate-001",
+        output_storage_uri: Optional[str] = None,
+        timeout: Optional[float] = None,
+    ) -> None:
+        r"""Initialize the Vertex AI Veo toolkit.
+
+        Args:
+            project_id (Optional[str]): Google Cloud project ID. If not
+                provided,
+                will use the default project from environment.
+                (default: :obj:`None`)
+            location (str): Google Cloud location for the API calls.
+                (default: :obj:`"us-central1"`)
+            model_id (str): The Veo model ID to use. Options include
+                "veo-2.0-generate-001" or "veo-3.0-generate-preview".
+                (default: :obj:`"veo-2.0-generate-001"`)
+            output_storage_uri (Optional[str]): Cloud Storage URI to save
+                output videos. If not provided, returns video bytes.
+                (default: :obj:`None`)
+            timeout (Optional[float]): Request timeout in seconds.
+                (default: :obj:`None`)
+        """
+        super().__init__(timeout)
+
+        from google.cloud import aiplatform
+
+        self.project_id = project_id or os.getenv("GOOGLE_CLOUD_PROJECT")
+        self.location = location
+        self.model_id = model_id
+        self.output_storage_uri = output_storage_uri
+
+        if not self.project_id:
+            raise ValueError(
+                "Project ID must be provided either as parameter or "
+                "GOOGLE_CLOUD_PROJECT environment variable"
+            )
+
+        aiplatform.init(project=self.project_id, location=self.location)
+        self.client = aiplatform.gapic.PredictionServiceClient()
+
+    def generate_video_from_text(
+        self,
+        text_prompt: str,
+        response_count: int = 1,
+        duration: int = 5,
+        aspect_ratio: str = "16:9",
+        negative_prompt: Optional[str] = None,
+        person_generation: str = "allow_adult",
+    ) -> Dict[str, Any]:
+        r"""Generate video from text prompt using Vertex AI Veo.
+
+        Args:
+            text_prompt (str): The text prompt to guide video generation.
+            response_count (int): Number of videos to generate (1-4).
+                (default: :obj:`1`)
+            duration (int): Video duration in seconds (5-8).
+                (default: :obj:`5`)
+            aspect_ratio (str): Video aspect ratio. Options: "16:9", "9:16".
+                (default: :obj:`"16:9"`)
+            negative_prompt (Optional[str]): What to avoid in the video.
+                (default: :obj:`None`)
+            person_generation (str): Person safety setting. Options:
+                "allow_adult", "dont_allow". (default: :obj:`"allow_adult"`)
+
+        Returns:
+            Dict[str, Any]:
+                A dictionary containing:
+                - 'success' (bool): Whether the operation was successful
+                - 'videos' (List[Dict]): List of generated video data
+                - 'metadata' (Dict): Additional metadata from the response
+                - 'error' (str): Error message if operation failed
+        """
+        try:
+            from google.protobuf import json_format
+            from google.protobuf.struct_pb2 import Value
+
+            # Construct the request
+            endpoint = (
+                f"projects/{self.project_id}/locations/{self.location}/"
+                f"publishers/google/models/{self.model_id}"
+            )
+
+            # Build parameters
+            parameters = {
+                "aspectRatio": aspect_ratio,
+                "personGeneration": person_generation,
+            }
+
+            if negative_prompt:
+                parameters["negativePrompt"] = negative_prompt
+
+            # Build request body
+            request_body = {
+                "contents": [
+                    {"role": "user", "parts": [{"text": text_prompt}]}
+                ],
+                "generationConfig": {
+                    "responseCount": response_count,
+                    "duration": duration,
+                },
+                "parameters": parameters,
+            }
+
+            if self.output_storage_uri:
+                generation_config = request_body.setdefault(
+                    "generationConfig", {}
+                )
+                if isinstance(generation_config, dict):
+                    generation_config["outputStorageUri"] = (
+                        self.output_storage_uri
+                    )
+
+            # Convert to protobuf format
+            request_value = json_format.ParseDict(request_body, Value())
+
+            # Make the API call
+            response = self.client.predict(
+                endpoint=endpoint,
+                instances=[request_value],
+                timeout=self.timeout,
+            )
+
+            return self._parse_video_response(response)
+
+        except Exception as e:
+            logger.error(f"Error generating video from text: {e}")
+            return {"error": str(e)}
+
+    def generate_video_from_image(
+        self,
+        image_path: str,
+        text_prompt: str,
+        response_count: int = 1,
+        duration: int = 5,
+        aspect_ratio: str = "16:9",
+        negative_prompt: Optional[str] = None,
+        person_generation: str = "allow_adult",
+    ) -> Dict[str, Any]:
+        r"""Generate video from image and text prompt using Vertex AI Veo.
+
+        Args:
+            image_path (str): Path to the input image file (local or GCS URI).
+            text_prompt (str): The text prompt to guide video generation.
+            response_count (int): Number of videos to generate (1-4).
+                (default: :obj:`1`)
+            duration (int): Video duration in seconds (5-8).
+                (default: :obj:`5`)
+            aspect_ratio (str): Video aspect ratio. Options: "16:9", "9:16".
+                (default: :obj:`"16:9"`)
+            negative_prompt (Optional[str]): What to avoid in the video.
+                (default: :obj:`None`)
+            person_generation (str): Person safety setting.
+                (default: :obj:`"allow_adult"`)
+
+        Returns:
+            Dict[str, Any]:
+                A dictionary containing:
+                - 'success' (bool): Whether the operation was successful
+                - 'videos' (List[Dict]): List of generated video data
+                - 'metadata' (Dict): Additional metadata from the response
+                - 'error' (str): Error message if operation failed
+        """
+        try:
+            from google.protobuf import json_format
+            from google.protobuf.struct_pb2 import Value
+
+            # Read and encode image
+            image_data, mime_type = self._process_image(image_path)
+
+            # Construct the request
+            endpoint = (
+                f"projects/{self.project_id}/locations/{self.location}/"
+                f"publishers/google/models/{self.model_id}"
+            )
+
+            # Build parameters
+            parameters = {
+                "aspectRatio": aspect_ratio,
+                "personGeneration": person_generation,
+            }
+
+            if negative_prompt:
+                parameters["negativePrompt"] = negative_prompt
+
+            # Build request body with image
+            request_body = {
+                "contents": [
+                    {
+                        "role": "user",
+                        "parts": [
+                            {
+                                "inlineData": {
+                                    "mimeType": mime_type,
+                                    "data": image_data,
+                                }
+                            },
+                            {"text": text_prompt},
+                        ],
+                    }
+                ],
+                "generationConfig": {
+                    "responseCount": response_count,
+                    "duration": duration,
+                },
+                "parameters": parameters,
+            }
+
+            if self.output_storage_uri:
+                generation_config = request_body.setdefault(
+                    "generationConfig", {}
+                )
+                if isinstance(generation_config, dict):
+                    generation_config["outputStorageUri"] = (
+                        self.output_storage_uri
+                    )
+
+            # Convert to protobuf format
+            request_value = json_format.ParseDict(request_body, Value())
+
+            # Make the API call
+            response = self.client.predict(
+                endpoint=endpoint,
+                instances=[request_value],
+                timeout=self.timeout,
+            )
+
+            return self._parse_video_response(response)
+
+        except Exception as e:
+            logger.error(f"Error generating video from image: {e}")
+            return {"error": str(e)}
+
+    def extend_video(
+        self,
+        video_uri: str,
+        text_prompt: str,
+        duration: int = 5,
+        aspect_ratio: str = "16:9",
+        negative_prompt: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        r"""Extend an existing video using Vertex AI Veo.
+
+        Args:
+            video_uri (str): Cloud Storage URI of the video to extend.
+            text_prompt (str): The text prompt to guide video extension.
+            duration (int): Duration to extend in seconds (5-8).
+                (default: :obj:`5`)
+            aspect_ratio (str): Video aspect ratio.
+                (default: :obj:`"16:9"`)
+            negative_prompt (Optional[str]): What to avoid in the extension.
+                (default: :obj:`None`)
+
+        Returns:
+            Dict[str, Any]:
+                A dictionary containing:
+                - 'success' (bool): Whether the operation was successful
+                - 'videos' (List[Dict]): List of extended video data
+                - 'metadata' (Dict): Additional metadata from the response
+                - 'error' (str): Error message if operation failed
+        """
+        try:
+            from google.protobuf import json_format
+            from google.protobuf.struct_pb2 import Value
+
+            # Construct the request
+            endpoint = (
+                f"projects/{self.project_id}/locations/{self.location}/"
+                f"publishers/google/models/{self.model_id}"
+            )
+
+            # Build parameters
+            parameters = {
+                "aspectRatio": aspect_ratio,
+                "videoToExtend": video_uri,
+            }
+
+            if negative_prompt:
+                parameters["negativePrompt"] = negative_prompt
+
+            # Build request body
+            request_body = {
+                "contents": [
+                    {"role": "user", "parts": [{"text": text_prompt}]}
+                ],
+                "generationConfig": {
+                    "duration": duration,
+                },
+                "parameters": parameters,
+            }
+
+            if self.output_storage_uri:
+                generation_config = request_body.setdefault(
+                    "generationConfig", {}
+                )
+                if isinstance(generation_config, dict):
+                    generation_config["outputStorageUri"] = (
+                        self.output_storage_uri
+                    )
+
+            # Convert to protobuf format
+            request_value = json_format.ParseDict(request_body, Value())
+
+            # Make the API call
+            response = self.client.predict(
+                endpoint=endpoint,
+                instances=[request_value],
+                timeout=self.timeout,
+            )
+
+            return self._parse_video_response(response)
+
+        except Exception as e:
+            logger.error(f"Error extending video: {e}")
+            return {"error": str(e)}
+
+    async def agenerate_video_from_text(
+        self,
+        text_prompt: str,
+        response_count: int = 1,
+        duration: int = 5,
+        aspect_ratio: str = "16:9",
+        negative_prompt: Optional[str] = None,
+        person_generation: str = "allow_adult",
+    ) -> Dict[str, Any]:
+        r"""Generate video from text prompt using Vertex AI Veo.
+
+        Args:
+            text_prompt (str): The text prompt to guide video generation.
+            response_count (int): Number of videos to generate (1-4).
+                (default: :obj:`1`)
+            duration (int): Video duration in seconds (5-8).
+                (default: :obj:`5`)
+            aspect_ratio (str): Video aspect ratio. Options: "16:9", "9:16".
+                (default: :obj:`"16:9"`)
+            negative_prompt (Optional[str]): What to avoid in the video.
+                (default: :obj:`None`)
+            person_generation (str): Person safety setting. Options:
+                "allow_adult", "dont_allow". (default: :obj:`"allow_adult"`)
+
+        Returns:
+            Dict[str, Any]:
+                A dictionary containing:
+                - 'success' (bool): Whether the operation was successful
+                - 'videos' (List[Dict]): List of generated video data
+                - 'metadata' (Dict): Additional metadata from the response
+                - 'error' (str): Error message if operation failed
+        """
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None,
+            self.generate_video_from_text,
+            text_prompt,
+            response_count,
+            duration,
+            aspect_ratio,
+            negative_prompt,
+            person_generation,
+        )
+
+    async def agenerate_video_from_image(
+        self,
+        image_path: str,
+        text_prompt: str,
+        response_count: int = 1,
+        duration: int = 5,
+        aspect_ratio: str = "16:9",
+        negative_prompt: Optional[str] = None,
+        person_generation: str = "allow_adult",
+    ) -> Dict[str, Any]:
+        r"""Generate video from image and text prompt using Vertex AI Veo.
+
+        Args:
+            image_path (str): Path to the input image file (local or GCS URI).
+            text_prompt (str): The text prompt to guide video generation.
+            response_count (int): Number of videos to generate (1-4).
+                (default: :obj:`1`)
+            duration (int): Video duration in seconds (5-8).
+                (default: :obj:`5`)
+            aspect_ratio (str): Video aspect ratio. Options: "16:9", "9:16".
+                (default: :obj:`"16:9"`)
+            negative_prompt (Optional[str]): What to avoid in the video.
+                (default: :obj:`None`)
+            person_generation (str): Person safety setting.
+                (default: :obj:`"allow_adult"`)
+
+        Returns:
+            Dict[str, Any]:
+                A dictionary containing:
+                - 'success' (bool): Whether the operation was successful
+                - 'videos' (List[Dict]): List of generated video data
+                - 'metadata' (Dict): Additional metadata from the response
+                - 'error' (str): Error message if operation failed
+        """
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None,
+            self.generate_video_from_image,
+            image_path,
+            text_prompt,
+            response_count,
+            duration,
+            aspect_ratio,
+            negative_prompt,
+            person_generation,
+        )
+
+    async def aextend_video(
+        self,
+        video_uri: str,
+        text_prompt: str,
+        duration: int = 5,
+        aspect_ratio: str = "16:9",
+        negative_prompt: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        r"""Extend an existing video using Vertex AI Veo.
+
+        Args:
+            video_uri (str): Cloud Storage URI of the video to extend.
+            text_prompt (str): The text prompt to guide video extension.
+            duration (int): Duration to extend in seconds (5-8).
+                (default: :obj:`5`)
+            aspect_ratio (str): Video aspect ratio.
+                (default: :obj:`"16:9"`)
+            negative_prompt (Optional[str]): What to avoid in the extension.
+                (default: :obj:`None`)
+
+        Returns:
+            Dict[str, Any]:
+                A dictionary containing:
+                - 'success' (bool): Whether the operation was successful
+                - 'videos' (List[Dict]): List of extended video data
+                - 'metadata' (Dict): Additional metadata from the response
+                - 'error' (str): Error message if operation failed
+        """
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None,
+            self.extend_video,
+            video_uri,
+            text_prompt,
+            duration,
+            aspect_ratio,
+            negative_prompt,
+        )
+
+    def _process_image(self, image_path: str) -> Tuple[str, str]:
+        r"""Process image file and return base64 encoded data and MIME type."""
+        if image_path.startswith('gs://'):
+            # Handle Google Cloud Storage URIs
+            from google.cloud import storage  # type: ignore[attr-defined]
+
+            # Parse GCS URI
+            parts = image_path[5:].split('/', 1)
+            bucket_name = parts[0]
+            blob_name = parts[1]
+
+            # Download image from GCS
+            client = storage.Client()
+            bucket = client.bucket(bucket_name)
+            blob = bucket.blob(blob_name)
+            image_bytes = blob.download_as_bytes()
+
+            # Determine MIME type from file extension
+            if blob_name.lower().endswith('.png'):
+                mime_type = "image/png"
+            elif blob_name.lower().endswith(('.jpg', '.jpeg')):
+                mime_type = "image/jpeg"
+            else:
+                raise ValueError("Unsupported image format. Use PNG or JPEG.")
+
+        else:
+            # Handle local file paths
+            with open(image_path, 'rb') as f:
+                image_bytes = f.read()
+
+            # Determine MIME type from file extension
+            if image_path.lower().endswith('.png'):
+                mime_type = "image/png"
+            elif image_path.lower().endswith(('.jpg', '.jpeg')):
+                mime_type = "image/jpeg"
+            else:
+                raise ValueError("Unsupported image format. Use PNG or JPEG.")
+
+        # Encode to base64
+        image_data = base64.b64encode(image_bytes).decode('utf-8')
+
+        return image_data, mime_type
+
+    def _parse_video_response(self, response: Any) -> Dict[str, Any]:
+        r"""Parse the video generation response."""
+        try:
+            result: Dict[str, Any] = {
+                "success": True,
+                "videos": [],
+                "metadata": {},
+            }
+
+            # Extract prediction results
+            if hasattr(response, 'predictions'):
+                for prediction in response.predictions:
+                    # Convert prediction to dict if needed
+                    if hasattr(prediction, 'struct_value'):
+                        pred_dict = dict(prediction.struct_value)
+                    else:
+                        pred_dict = prediction
+
+                    videos_list = result["videos"]
+                    if isinstance(videos_list, list):
+                        videos_list.append(pred_dict)
+
+            # Extract metadata if available
+            if hasattr(response, 'metadata'):
+                result["metadata"] = dict(response.metadata)
+
+            return result
+
+        except Exception as e:
+            logger.error(f"Error parsing video response: {e}")
+            return {
+                "success": False,
+                "error": str(e),
+                "raw_response": str(response),
+            }
+
+    def get_tools(self) -> List[FunctionTool]:
+        r"""Get a list of tools for video generation with Vertex AI Veo.
+
+        Returns:
+            List[FunctionTool]: List of available function tools.
+        """
+        return [
+            FunctionTool(self.generate_video_from_text),
+            FunctionTool(self.generate_video_from_image),
+            FunctionTool(self.extend_video),
+        ]
+
+    def get_async_tools(self) -> List[FunctionTool]:
+        r"""Get a list of async tools for video generation with Vertex AI Veo.
+
+        Returns:
+            List[FunctionTool]: List of available async function tools.
+        """
+        return [
+            FunctionTool(self.agenerate_video_from_text),
+            FunctionTool(self.agenerate_video_from_image),
+            FunctionTool(self.aextend_video),
+        ]