feat(xai): add Grok Imagine for image and video generation

.gitignore

@@ -159,3 +159,4 @@ uv.lock
 bandit-report.json
 
 scripts/
+assets/

notebooks/working-with-images.ipynb

@@ -20,13 +20,13 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "import celeste\n",
     "from IPython.display import Image, display"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
@@ -41,24 +41,24 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "img_gen_result = await celeste.images.generate(\n",
     "    \"A nano banana on the beach\",\n",
     "    model=\"gemini-2.5-flash-image\",\n",
     ")"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "display(Image(data=img_gen_result.content.data))"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
@@ -73,25 +73,25 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "img_edit_result = await celeste.images.edit(\n",
     "    image=img_gen_result.content,\n",
     "    prompt=\"Make it night time\",\n",
     "    model=\"gemini-2.5-flash-image\",\n",
     ")"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "display(Image(data=img_edit_result.content.data))"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
@@ -106,25 +106,25 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "analyze_result = await celeste.images.analyze(\n",
     "    prompt=\"What fruit is in this image and what color is it?\",\n",
     "    image=img_gen_result.content,\n",
     "    model=\"gemini-2.5-flash-lite\",\n",
     ")"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "print(analyze_result.content)"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
@@ -159,7 +159,9 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "prompt = \"A blurry iPhone-style photograph showing the window of a moving train. Through the window, a scenic landscape appears: tall green cliffs running alongside a river, with a small European village built on the slopes. The motion blur suggests the train is moving quickly, with soft reflections on the glass, natural daylight, and a casual handheld phone-camera aesthetic. Sharp textures where possible, rich colors, and a realistic sense of depth and distance.\"\n",
     "\n",
@@ -170,13 +172,11 @@
     "    steps=1,\n",
     ")\n",
     "display(Image(data=local_result.content.data))"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
-   "metadata": {},
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "---\n",
     "\n",
@@ -187,7 +187,9 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "from tqdm.asyncio import tqdm\n",
     "\n",
@@ -204,24 +206,22 @@
     "    pass\n",
     "\n",
     "display(Image(data=chunk.content.data))"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
-   "metadata": {},
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "---\n",
     "Star on GitHub 👉 [withceleste/celeste-python](https://github.com/withceleste/celeste-python)"
    ]
   },
   {
-   "metadata": {},
    "cell_type": "code",
-   "outputs": [],
    "execution_count": null,
-   "source": ""
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "celeste-ai"
-version = "0.9.5"
+version = "0.9.6"
 description = "Open source, type-safe primitives for multi-modal AI. All capabilities, all providers, one interface"
 authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}]
 readme = "README.md"

src/celeste/modalities/images/models.py

@@ -6,10 +6,12 @@
 from .providers.byteplus.models import MODELS as BYTEPLUS_MODELS
 from .providers.google.models import MODELS as GOOGLE_MODELS
 from .providers.openai.models import MODELS as OPENAI_MODELS
+from .providers.xai.models import MODELS as XAI_MODELS
 
 MODELS: list[Model] = [
     *BFL_MODELS,
     *BYTEPLUS_MODELS,
     *GOOGLE_MODELS,
     *OPENAI_MODELS,
+    *XAI_MODELS,
 ]

src/celeste/modalities/images/providers/__init__.py

@@ -8,11 +8,13 @@
 from .google import GoogleImagesClient
 from .ollama import OllamaImagesClient
 from .openai import OpenAIImagesClient
+from .xai import XAIImagesClient
 
 PROVIDERS: dict[Provider, type[ImagesClient]] = {
     Provider.BFL: BFLImagesClient,
     Provider.BYTEPLUS: BytePlusImagesClient,
     Provider.GOOGLE: GoogleImagesClient,
     Provider.OLLAMA: OllamaImagesClient,
     Provider.OPENAI: OpenAIImagesClient,
+    Provider.XAI: XAIImagesClient,
 }

src/celeste/modalities/images/providers/xai/__init__.py

@@ -0,0 +1,6 @@
+"""xAI provider for images modality."""
+
+from .client import XAIImagesClient
+from .models import MODELS
+
+__all__ = ["MODELS", "XAIImagesClient"]

src/celeste/modalities/images/providers/xai/client.py

@@ -0,0 +1,101 @@
+"""xAI images client."""
+
+from typing import Any, Unpack
+
+from celeste.artifacts import ImageArtifact
+from celeste.parameters import ParameterMapper
+from celeste.providers.xai.images import config
+from celeste.providers.xai.images.client import XAIImagesClient as XAIImagesMixin
+
+from ...client import ImagesClient
+from ...io import (
+    ImageFinishReason,
+    ImageInput,
+    ImageOutput,
+    ImageUsage,
+)
+from ...parameters import ImageParameters
+from .parameters import XAI_PARAMETER_MAPPERS
+
+
+class XAIImagesClient(XAIImagesMixin, ImagesClient):
+    """xAI images client."""
+
+    @classmethod
+    def parameter_mappers(cls) -> list[ParameterMapper]:
+        return XAI_PARAMETER_MAPPERS
+
+    def _init_request(self, inputs: ImageInput) -> dict[str, Any]:
+        """Initialize request from inputs."""
+        request: dict[str, Any] = {"prompt": inputs.prompt}
+        if inputs.image is not None:
+            # xAI accepts URL or base64 string
+            if inputs.image.url:
+                request["image"] = inputs.image.url
+            else:
+                request["image"] = inputs.image.get_base64()
+        return request
+
+    async def generate(
+        self,
+        prompt: str,
+        **parameters: Unpack[ImageParameters],
+    ) -> ImageOutput:
+        """Generate images from prompt."""
+        inputs = ImageInput(prompt=prompt)
+        return await self._predict(
+            inputs,
+            endpoint=config.XAIImagesEndpoint.CREATE_IMAGE,
+            **parameters,
+        )
+
+    async def edit(
+        self,
+        image: ImageArtifact,
+        prompt: str,
+        **parameters: Unpack[ImageParameters],
+    ) -> ImageOutput:
+        """Edit an image with text instructions."""
+        inputs = ImageInput(image=image, prompt=prompt)
+        return await self._predict(
+            inputs,
+            endpoint=config.XAIImagesEndpoint.CREATE_EDIT,
+            **parameters,
+        )
+
+    def _parse_usage(self, response_data: dict[str, Any]) -> ImageUsage:
+        """Parse usage from response."""
+        usage = super()._parse_usage(response_data)
+        return ImageUsage(**usage)
+
+    def _parse_content(
+        self,
+        response_data: dict[str, Any],
+        **parameters: Unpack[ImageParameters],
+    ) -> ImageArtifact:
+        """Parse content from response."""
+        data = super()._parse_content(response_data)
+        image_data = data[0]
+
+        # xAI returns either b64_json or url
+        b64_json = image_data.get("b64_json")
+        if b64_json:
+            import base64
+
+            image_bytes = base64.b64decode(b64_json)
+            return ImageArtifact(data=image_bytes)
+
+        url = image_data.get("url")
+        if url:
+            return ImageArtifact(url=url)
+
+        msg = "No image URL or base64 data in response"
+        raise ValueError(msg)
+
+    def _parse_finish_reason(self, response_data: dict[str, Any]) -> ImageFinishReason:
+        """Parse finish reason from response."""
+        finish_reason = super()._parse_finish_reason(response_data)
+        return ImageFinishReason(reason=finish_reason.reason)
+
+
+__all__ = ["XAIImagesClient"]

src/celeste/modalities/images/providers/xai/models.py

@@ -0,0 +1,38 @@
+"""xAI models for images modality."""
+
+from celeste.constraints import Choice, Range
+from celeste.core import Modality, Operation, Provider
+from celeste.models import Model
+
+from ...parameters import ImageParameter
+
+MODELS: list[Model] = [
+    Model(
+        id="grok-imagine-image",
+        provider=Provider.XAI,
+        display_name="Grok Imagine Image",
+        operations={Modality.IMAGES: {Operation.GENERATE, Operation.EDIT}},
+        parameter_constraints={
+            ImageParameter.NUM_IMAGES: Range(min=1, max=10),
+            ImageParameter.ASPECT_RATIO: Choice(
+                options=[
+                    "1:1",
+                    "3:4",
+                    "4:3",
+                    "9:16",
+                    "16:9",
+                    "2:3",
+                    "3:2",
+                    "9:19.5",
+                    "19.5:9",
+                    "9:20",
+                    "20:9",
+                    "1:2",
+                    "2:1",
+                    "auto",
+                ]
+            ),
+            ImageParameter.OUTPUT_FORMAT: Choice(options=["url", "b64_json"]),
+        },
+    ),
+]