feat: Implement embed_image() (#5101)

## Changes Made

Adds support for `embed_image()`, e.g.

```
import daft
from daft.functions.ai import embed_image

import numpy as np

provider = "transformers"
model = "openai/clip-vit-base-patch32"
test_image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)

(
    daft.from_pydict({"image": [test_image] * 16})
    .select(daft.col("image").cast(daft.DataType.image()))
    .select(embed_image(daft.col("image"), provider=provider, model=model))
    .show()
)
```

**!! Currently only supports OpenAI CLIP models:
https://huggingface.co/docs/transformers/en/model_doc/clip**

---------

Co-authored-by: R. C. Howell <[email protected]>

Author: desmondcheongzx

daft/ai/_expressions.py

@@ -4,7 +4,7 @@
 
 if TYPE_CHECKING:
     from daft import Series
-    from daft.ai.protocols import TextEmbedder, TextEmbedderDescriptor
+    from daft.ai.protocols import ImageEmbedder, ImageEmbedderDescriptor, TextEmbedder, TextEmbedderDescriptor
     from daft.ai.typing import Embedding
 
 
@@ -19,3 +19,16 @@ def __init__(self, text_embedder: TextEmbedderDescriptor):
     def __call__(self, text_series: Series) -> list[Embedding]:
         text = text_series.to_pylist()
         return self.text_embedder.embed_text(text) if text else []
+
+
+class _ImageEmbedderExpression:
+    """Function expression implementation for an ImageEmbedder protocol."""
+
+    image_embedder: ImageEmbedder
+
+    def __init__(self, image_embedder: ImageEmbedderDescriptor):
+        self.image_embedder = image_embedder.instantiate()
+
+    def __call__(self, image_series: Series) -> list[Embedding]:
+        image = image_series.to_pylist()
+        return self.image_embedder.embed_image(image) if image else []

daft/ai/openai/__init__.py

@@ -9,7 +9,7 @@
 
 if TYPE_CHECKING:
     from daft.ai.openai.typing import OpenAIProviderOptions
-    from daft.ai.protocols import TextEmbedder, TextEmbedderDescriptor
+    from daft.ai.protocols import ImageEmbedderDescriptor, TextEmbedderDescriptor
     from daft.ai.typing import Options
 
 __all__ = [
@@ -36,3 +36,6 @@ def get_text_embedder(self, model: str | None = None, **options: Any) -> TextEmb
             model_name=(model or "text-embedding-3-small"),
             model_options=options,
         )
+
+    def get_image_embedder(self, model: str | None = None, **options: Any) -> ImageEmbedderDescriptor:
+        raise NotImplementedError("embed_image is not currently implemented for the OpenAI provider")

daft/ai/protocols.py

@@ -6,7 +6,7 @@
 from daft.ai.typing import Descriptor
 
 if TYPE_CHECKING:
-    from daft.ai.typing import Embedding, EmbeddingDimensions
+    from daft.ai.typing import Embedding, EmbeddingDimensions, Image
 
 
 @runtime_checkable
@@ -24,3 +24,20 @@ class TextEmbedderDescriptor(Descriptor[TextEmbedder]):
     @abstractmethod
     def get_dimensions(self) -> EmbeddingDimensions:
         """Returns the dimensions of the embeddings produced by the described TextEmbedder."""
+
+
+@runtime_checkable
+class ImageEmbedder(Protocol):
+    """Protocol for image embedding implementations."""
+
+    def embed_image(self, images: list[Image]) -> list[Embedding]:
+        """Embeds a batch of images into an embedding vector."""
+        ...
+
+
+class ImageEmbedderDescriptor(Descriptor[ImageEmbedder]):
+    """Descriptor for a ImageEmbedder implementation."""
+
+    @abstractmethod
+    def get_dimensions(self) -> EmbeddingDimensions:
+        """Returns the dimensions of the embeddings produced by the described ImageEmbedder."""

daft/ai/provider.py

@@ -7,7 +7,7 @@
 
 if TYPE_CHECKING:
     from daft.ai.openai.typing import OpenAIProviderOptions
-    from daft.ai.protocols import TextEmbedderDescriptor
+    from daft.ai.protocols import ImageEmbedderDescriptor, TextEmbedderDescriptor
 
 
 class ProviderImportError(ImportError):
@@ -34,9 +34,19 @@ def load_sentence_transformers(name: str | None = None, **options: Any) -> Provi
         raise ProviderImportError(["sentence_transformers", "torch"]) from e
 
 
+def load_transformers(name: str | None = None, **options: Any) -> Provider:
+    try:
+        from daft.ai.transformers import TransformersProvider
+
+        return TransformersProvider(name, **options)
+    except ImportError as e:
+        raise ProviderImportError(["torch", "torchvision", "transformers", "Pillow"]) from e
+
+
 PROVIDERS: dict[str, Callable[..., Provider]] = {
     "openai": load_openai,
     "sentence_transformers": load_sentence_transformers,
+    "transformers": load_transformers,
 }
 
 
@@ -65,3 +75,8 @@ def name(self) -> str:
     def get_text_embedder(self, model: str | None = None, **options: Any) -> TextEmbedderDescriptor:
         """Returns a TextEmbedderDescriptor for this provider."""
         ...
+
+    @abstractmethod
+    def get_image_embedder(self, model: str | None = None, **options: Any) -> ImageEmbedderDescriptor:
+        """Returns an ImageEmbedderDescriptor for this provider."""
+        ...

daft/ai/sentence_transformers/__init__.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
-    from daft.ai.protocols import TextEmbedderDescriptor
+    from daft.ai.protocols import ImageEmbedderDescriptor, TextEmbedderDescriptor
     from daft.ai.typing import Options
 
 __all__ = [
@@ -28,3 +28,6 @@ def name(self) -> str:
 
     def get_text_embedder(self, model: str | None = None, **options: Any) -> TextEmbedderDescriptor:
         return SentenceTransformersTextEmbedderDescriptor(model or "sentence-transformers/all-MiniLM-L6-v2", options)
+
+    def get_image_embedder(self, model: str | None = None, **options: Any) -> ImageEmbedderDescriptor:
+        raise NotImplementedError("embed_image is not currently implemented for the Sentence Transformers provider")

daft/ai/sentence_transformers/text_embedder.py

@@ -30,7 +30,7 @@ def get_options(self) -> Options:
         return self.options
 
     def get_dimensions(self) -> EmbeddingDimensions:
-        dimensions = AutoConfig.from_pretrained(self.model).hidden_size
+        dimensions = AutoConfig.from_pretrained(self.model, trust_remote_code=True).hidden_size
         return EmbeddingDimensions(size=dimensions, dtype=DataType.float32())
 
     def instantiate(self) -> TextEmbedder:

daft/ai/transformers/__init__.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from daft.ai.provider import Provider
+
+from daft.ai.transformers.image_embedder import TransformersImageEmbedderDescriptor
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from daft.ai.protocols import ImageEmbedderDescriptor, TextEmbedderDescriptor
+    from daft.ai.typing import Options
+
+__all__ = [
+    "TransformersProvider",
+]
+
+
+class TransformersProvider(Provider):
+    _name: str
+    _options: Options
+
+    def __init__(self, name: str | None = None, **options: Any):
+        self._name = name if name else "transformers"
+        self._options = options
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def get_image_embedder(self, model: str | None = None, **options: Any) -> ImageEmbedderDescriptor:
+        return TransformersImageEmbedderDescriptor(model or "openai/clip-vit-base-patch32", options)
+
+    def get_text_embedder(self, model: str | None = None, **options: Any) -> TextEmbedderDescriptor:
+        raise NotImplementedError("embed_text is not currently implemented for the Transformers provider")

daft/ai/transformers/image_embedder.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import torch
+from transformers import AutoConfig, AutoModel, AutoProcessor
+
+from daft import DataType
+from daft.ai.protocols import ImageEmbedder, ImageEmbedderDescriptor
+from daft.ai.typing import EmbeddingDimensions, Options
+from daft.ai.utils import get_device
+from daft.dependencies import pil_image
+
+if TYPE_CHECKING:
+    from daft.ai.typing import Embedding, Image
+
+
+@dataclass
+class TransformersImageEmbedderDescriptor(ImageEmbedderDescriptor):
+    model: str
+    options: Options
+
+    def get_provider(self) -> str:
+        return "transformers"
+
+    def get_model(self) -> str:
+        return self.model
+
+    def get_options(self) -> Options:
+        return self.options
+
+    def get_dimensions(self) -> EmbeddingDimensions:
+        config = AutoConfig.from_pretrained(self.model, trust_remote_code=True)
+        # For CLIP models, the image embedding dimension is typically in projection_dim or hidden_size.
+        embedding_size = getattr(config, "projection_dim", getattr(config, "hidden_size", 512))
+        return EmbeddingDimensions(size=embedding_size, dtype=DataType.float32())
+
+    def instantiate(self) -> ImageEmbedder:
+        return TransformersImageEmbedder(self.model, **self.options)
+
+
+class TransformersImageEmbedder(ImageEmbedder):
+    model: Any
+    options: Options
+
+    def __init__(self, model_name_or_path: str, **options: Any):
+        self.device = get_device()
+        self.model = AutoModel.from_pretrained(
+            model_name_or_path,
+            trust_remote_code=True,
+            use_safetensors=True,
+        ).to(self.device)
+        self.processor = AutoProcessor.from_pretrained(model_name_or_path, trust_remote_code=True, use_fast=True)
+        self.options = options
+
+    def embed_image(self, images: list[Image]) -> list[Embedding]:
+        # TODO(desmond): There's potential for image decoding and processing on the GPU with greater
+        # performance. Methods differ a little between different models, so let's do it later.
+        pil_images = [pil_image.fromarray(image) for image in images]
+        processed = self.processor(images=pil_images, return_tensors="pt")
+        pixel_values = processed["pixel_values"].to(self.device)
+
+        with torch.inference_mode():
+            embeddings = self.model.get_image_features(pixel_values)
+        return embeddings.cpu().numpy().tolist()

daft/ai/typing.py

@@ -16,6 +16,7 @@
     "Descriptor",
     "Embedding",
     "EmbeddingDimensions",
+    "Image",
 ]
 
 
@@ -47,8 +48,10 @@ def instantiate(self) -> T:
     from daft.dependencies import np
 
     Embedding: TypeAlias = np.typing.NDArray[Any]
+    Image: TypeAlias = np.ndarray[Any, Any]
 else:
     Embedding: TypeAlias = Any
+    Image: TypeAlias = Any
 
 
 @dataclass

daft/ai/utils.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch
+
+
+def get_device() -> torch.device:
+    """Get the best available PyTorch device for computation.
+
+    This function automatically selects the optimal device in order of preference:
+    1. CUDA GPU (if available) - for NVIDIA GPUs with CUDA support
+    2. MPS (Metal Performance Shaders) - for Apple Silicon Macs
+    3. CPU - as fallback when no GPU acceleration is available
+    """
+    import torch
+
+    device = (
+        torch.device("cuda")
+        if torch.cuda.is_available()
+        else torch.device("mps")
+        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+        else torch.device("cpu")
+    )
+    return device