Feature/Luma Ai Video Generation Driver

griptape/artifacts/__init__.py

@@ -9,6 +9,7 @@
 from .image_artifact import ImageArtifact
 from .audio_artifact import AudioArtifact
 from .action_artifact import ActionArtifact
+from .video_artifact import VideoArtifact
 from .generic_artifact import GenericArtifact
 
 
@@ -24,5 +25,6 @@
     "ImageArtifact",
     "AudioArtifact",
     "ActionArtifact",
+    "VideoArtifact",
     "GenericArtifact",
 ]

griptape/artifacts/video_artifact.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from attrs import define, field
+
+from griptape.artifacts import BlobArtifact
+
+
+@define
+class VideoArtifact(BlobArtifact):
+    """Stores video binary data and relevant metadata.
+
+    Attributes:
+        value: The video binary data.
+        mime_type: The video MIME type.
+        resolution: The resolution of the video (e.g., 1920x1080).
+        duration: Duration of the video in seconds.
+    """
+
+    aspect_ratio: tuple[int, int] = field(default=(16, 9), kw_only=True)
+
+    @property
+    def mime_type(self) -> str:
+        return "video/mp4"  # Or make this flexible based on the video format
+
+    def get_aspect_ratio(self) -> tuple[int, int]:
+        return self.aspect_ratio
+
+    def to_text(self) -> str:
+        raise NotImplementedError("VideoArtifact cannot be converted to text.")

griptape/drivers/__init__.py

@@ -135,6 +135,9 @@
 from .observability.griptape_cloud_observability_driver import GriptapeCloudObservabilityDriver
 from .observability.datadog_observability_driver import DatadogObservabilityDriver
 
+from .video_generation.base_video_generation_driver import BaseVideoGenerationDriver
+from .video_generation.dream_machine_video_generation_driver import DreamMachineVideoGenerationDriver
+
 __all__ = [
     "BasePromptDriver",
     "OpenAiChatPromptDriver",
@@ -242,4 +245,6 @@
     "OpenTelemetryObservabilityDriver",
     "GriptapeCloudObservabilityDriver",
     "DatadogObservabilityDriver",
+    "DreamMachineVideoGenerationDriver",
+    "BaseVideoGenerationDriver",
 ]

griptape/drivers/video_generation/base_video_generation_driver.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from attrs import define
+
+from griptape.events import EventBus, FinishVideoGenerationEvent, StartVideoGenerationEvent
+from griptape.mixins.exponential_backoff_mixin import ExponentialBackoffMixin
+from griptape.mixins.serializable_mixin import SerializableMixin
+
+if TYPE_CHECKING:
+    from griptape.artifacts import VideoArtifact
+
+
+@define
+class BaseVideoGenerationDriver(SerializableMixin, ExponentialBackoffMixin, ABC):
+    def before_run(self, prompt: str) -> None:
+        EventBus.publish_event(StartVideoGenerationEvent(prompt=prompt))
+
+    def after_run(self) -> None:
+        EventBus.publish_event(FinishVideoGenerationEvent())
+
+    def run_text_to_video(self, prompt: str) -> VideoArtifact:
+        for attempt in self.retrying():
+            with attempt:
+                self.before_run(prompt)
+                result = self.try_text_to_video(prompt)
+                self.after_run()
+
+                return result
+        else:
+            raise Exception("Failed to run text to video generation")
+
+    @abstractmethod
+    def try_text_to_video(self, prompt: str) -> VideoArtifact: ...

griptape/drivers/video_generation/dream_machine_video_generation_driver.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+import logging
+import time
+from typing import TYPE_CHECKING, Any
+
+import requests
+from attrs import Factory, define, field
+
+from griptape.artifacts import VideoArtifact
+from griptape.drivers import BaseVideoGenerationDriver
+from griptape.utils import import_optional_dependency
+
+if TYPE_CHECKING:
+    from lumaai import LumaAI
+
+logger = logging.getLogger(__name__)
+
+
+@define
+class DreamMachineVideoGenerationDriver(BaseVideoGenerationDriver):
+    api_key: str = field(kw_only=True, metadata={"serializable": True})
+    client: LumaAI = field(
+        default=Factory(
+            lambda self: import_optional_dependency("lumaai").LumaAI(auth_token=self.api_key), takes_self=True
+        ),
+        kw_only=True,
+    )
+    params: dict[str, Any] = field(default={}, kw_only=True, metadata={"serializable": True})
+
+    def try_text_to_video(self, prompt: str) -> VideoArtifact:
+        response = self.client.generations.create(prompt=prompt, **self.params)
+        generation = response
+        status = generation.state
+        while status in ["dreaming", "queued"]:
+            time.sleep(5)
+            if not generation.id:
+                raise Exception("Generation ID not found in the response")
+
+            generation = self.client.generations.get(generation.id)
+            status = generation.state
+        if status == "completed":
+            video_url = generation.assets.video
+            if not video_url:
+                raise Exception("Video URL not found in the generation response")
+            video_binary = self._download_video(video_url)
+            return VideoArtifact(
+                value=video_binary,
+            )
+        else:
+            raise Exception(f"Video generation failed with status: {status}")
+
+    def _download_video(self, video_url: str) -> bytes:
+        response = requests.get(video_url)
+        response.raise_for_status()
+        return response.content

griptape/events/__init__.py

@@ -22,6 +22,9 @@
 from .base_audio_transcription_event import BaseAudioTranscriptionEvent
 from .start_audio_transcription_event import StartAudioTranscriptionEvent
 from .finish_audio_transcription_event import FinishAudioTranscriptionEvent
+from .base_video_generation_event import BaseVideoGenerationEvent
+from .start_video_generation_event import StartVideoGenerationEvent
+from .finish_video_generation_event import FinishVideoGenerationEvent
 from .event_bus import EventBus
 
 __all__ = [
@@ -49,5 +52,8 @@
     "BaseAudioTranscriptionEvent",
     "StartAudioTranscriptionEvent",
     "FinishAudioTranscriptionEvent",
+    "BaseVideoGenerationEvent",
+    "StartVideoGenerationEvent",
+    "FinishVideoGenerationEvent",
     "EventBus",
 ]

griptape/events/base_video_generation_event.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from abc import ABC
+
+from attrs import define
+
+from .base_media_generation_event import BaseMediaGenerationEvent
+
+
+@define
+class BaseVideoGenerationEvent(BaseMediaGenerationEvent, ABC): ...

griptape/events/finish_video_generation_event.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from attrs import define
+
+from .base_video_generation_event import BaseVideoGenerationEvent
+
+
+@define
+class FinishVideoGenerationEvent(BaseVideoGenerationEvent): ...

griptape/events/start_video_generation_event.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from attrs import define, field
+
+from .base_video_generation_event import BaseVideoGenerationEvent
+
+
+@define
+class StartVideoGenerationEvent(BaseVideoGenerationEvent):
+    prompt: str = field(kw_only=True, metadata={"serializable": True})

griptape/loaders/__init__.py

@@ -9,6 +9,7 @@
 from .email_loader import EmailLoader
 from .image_loader import ImageLoader
 from .audio_loader import AudioLoader
+from .video_loader import VideoLoader
 from .blob_loader import BlobLoader
 
 
@@ -24,5 +25,6 @@
     "EmailLoader",
     "ImageLoader",
     "AudioLoader",
+    "VideoLoader",
     "BlobLoader",
 ]

griptape/loaders/video_loader.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from io import BytesIO
+from typing import Optional, cast
+
+from attrs import define, field
+
+from griptape.artifacts import VideoArtifact
+from griptape.loaders import BaseLoader
+from griptape.utils import import_optional_dependency
+
+
+@define
+class VideoLoader(BaseLoader):
+    """Loads videos into video artifacts.
+
+    Attributes:
+        format: If provided, attempts to ensure video artifacts are in this format when loaded.
+                For example, when set to 'mp4', loading video.webm will return a VideoArtifact containing the video
+                    bytes in MP4 format.
+    """
+
+    format: Optional[str] = field(default=None, kw_only=True)
+
+    FORMAT_TO_MIME_TYPE = {
+        "mp4": "video/mp4",
+        "webm": "video/webm",
+        "ogg": "video/ogg",
+    }
+
+    def load(self, source: bytes, *args, **kwargs) -> VideoArtifact:
+        moviepy = import_optional_dependency("moviepy.editor")
+        video = moviepy.VideoFileClip(BytesIO(source))
+
+        # Normalize format only if requested.
+        if self.format is not None:
+            byte_stream = BytesIO()
+            video.write_videofile(byte_stream, codec="libx264", format=self.format)
+            video = moviepy.VideoFileClip(byte_stream)
+            source = byte_stream.getvalue()
+        return VideoArtifact(source, aspect_ratio=(video.size[0], video.size[1]))
+
+    def _get_mime_type(self, video_format: str | None) -> str:
+        if video_format is None:
+            raise ValueError("video_format is None")
+
+        if video_format.lower() not in self.FORMAT_TO_MIME_TYPE:
+            raise ValueError(f"Unsupported video format {video_format}")
+
+        return self.FORMAT_TO_MIME_TYPE[video_format.lower()]
+
+    def load_collection(self, sources: list[bytes], *args, **kwargs) -> dict[str, VideoArtifact]:
+        return cast(dict[str, VideoArtifact], super().load_collection(sources, *args, **kwargs))

griptape/mixins/artifact_file_output_mixin.py

@@ -7,7 +7,7 @@
 from attrs import Attribute, define, field
 
 if TYPE_CHECKING:
-    from griptape.artifacts import BaseArtifact
+    from griptape.artifacts import BaseArtifact, VideoArtifact
 
 
 @define(slots=False)
@@ -43,3 +43,16 @@
             os.makedirs(os.path.dirname(outfile), exist_ok=True)
 
         Path(outfile).write_bytes(artifact.to_bytes())
+
+    def save_video_artifact(self, artifact: VideoArtifact) -> None:
+        if self.output_file:
+            outfile = self.output_file
+        elif self.output_dir:
+            outfile = os.path.join(self.output_dir, artifact.name + ".mp4")
+        else:
+            raise ValueError("No output_file or output_dir specified.")
+
+        if os.path.dirname(outfile):
+            os.makedirs(os.path.dirname(outfile), exist_ok=True)
+
+        Path(outfile).write_bytes(artifact.to_bytes())

griptape/tasks/__init__.py

@@ -19,6 +19,8 @@
 from .text_to_speech_task import TextToSpeechTask
 from .structure_run_task import StructureRunTask
 from .audio_transcription_task import AudioTranscriptionTask
+from .base_video_generation_task import BaseVideoGenerationTask
+from .prompt_video_generation_task import PromptVideoGenerationTask
 
 __all__ = [
     "BaseTask",
@@ -42,4 +44,6 @@
     "TextToSpeechTask",
     "StructureRunTask",
     "AudioTranscriptionTask",
+    "BaseVideoGenerationTask",
+    "PromptVideoGenerationTask",
 ]

griptape/tasks/base_video_generation_task.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import logging
+import os
+from abc import ABC
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from attrs import define
+
+from griptape.configs import Defaults
+from griptape.loaders import VideoLoader
+from griptape.mixins.artifact_file_output_mixin import ArtifactFileOutputMixin
+from griptape.tasks import BaseTask
+
+if TYPE_CHECKING:
+    from griptape.artifacts import VideoArtifact
+
+logger = logging.getLogger(Defaults.logging_config.logger_name)
+
+
+@define
+class BaseVideoGenerationTask(ArtifactFileOutputMixin, BaseTask, ABC):
+    """Provides a base class for video generation-related tasks.
+
+    Attributes:
+        output_dir: If provided, the generated video will be written to disk in output_dir.
+        output_file: If provided, the generated video will be written to disk as output_file.
+    """
+
+    def _read_from_file(self, path: str) -> VideoArtifact:
+        logger.info("Reading video from %s", os.path.abspath(path))
+        return VideoLoader().load(Path(path).read_bytes())