Skip to content

Add Imagine file storage & file-ID inputs to image/video generation #160

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Omar-V2 merged 1 commit into from
Jun 8, 2026
Merged

Add Imagine file storage & file-ID inputs to image/video generation #160

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
# Changelog

## [Unreleased]
### Added
- **Imagine File Storage**: Image and video generation (sync and async) now accept a `storage_options` parameter to persist generated assets to the Files API. It takes a dict with a required `filename` and optional `expires_after` (an `int` in seconds or a `datetime.timedelta`) and `public_url` (`True` to create a public URL with default expiry, or `{"expires_after": <seconds>}` for an independent URL expiry). Image and video responses expose new `file_output`, `storage_error`, `public_url`, and `public_url_error` properties.
- **File-ID Inputs for Generation**: Image and video generation now accept Files API `file_id` references as inputs alongside URLs/base64 — `image_file_id` / `image_file_ids` for `image.sample()` / `image.sample_batch()`, and `image_file_id` / `video_file_id` / `reference_image_file_ids` for `video.generate()` / `video.extend()` (and the batch `prepare` helpers). URL and file-ID lists may be mixed in the same multi-image request (file IDs are sent first).

## [v1.14.0]
### Added
Expand Down
95 changes: 88 additions & 7 deletions src/xai_sdk/aio/image.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from opentelemetry.trace import SpanKind

from ..__about__ import __version__
from ..files import StorageOptions
from ..image import (
BaseClient,
BaseImageResponse,
Expand All @@ -14,7 +15,7 @@
_make_span_request_attributes,
_make_span_response_attributes,
)
from ..proto import batch_pb2
from ..proto import batch_pb2, image_pb2
from ..telemetry import get_tracer
from ..types import ImageGenerationModel

Expand All @@ -31,11 +32,14 @@ def prepare(
*,
batch_request_id: Optional[str] = None,
image_url: Optional[str] = None,
image_file_id: Optional[str] = None,
image_urls: Optional[Sequence[str]] = None,
image_file_ids: Optional[Sequence[str]] = None,
user: Optional[str] = None,
image_format: Optional[ImageFormat] = None,
aspect_ratio: Optional[ImageAspectRatio] = None,
resolution: Optional[ImageResolution] = None,
storage_options: Optional[Union[StorageOptions, image_pb2.StorageOptions]] = None,
) -> batch_pb2.BatchRequest:
"""Prepares an image generation request for batch processing.

Expand All @@ -49,13 +53,35 @@ def prepare(
**If provided, it must be unique within the batch.** Used to identify the
corresponding result when the response is returned.
image_url: The URL or base64-encoded string of an input image to use as a starting point.
Cannot be set together with `image_urls`. Only supported for grok-imagine models.
Cannot be set together with `image_file_id`, `image_urls`, or `image_file_ids`.
Only supported for grok-imagine models.
image_file_id: The file ID of an input image to use as a starting point.
Cannot be set together with `image_url`, `image_urls`, or `image_file_ids`.
Only supported for grok-imagine models.
image_urls: Optional list of input images for multi-reference image editing.
Cannot be set together with `image_url`. Only supported for grok-imagine models.
Cannot be set together with single-image params (`image_url` /
`image_file_id`). May be combined with `image_file_ids` to mix
URL/base64 and file-ID references in the same request; file IDs
are appended first in that case.
Only supported for grok-imagine models.
image_file_ids: Optional list of input image file IDs for multi-reference image editing.
Cannot be set together with single-image params (`image_url` /
`image_file_id`). May be combined with `image_urls` to mix
URL/base64 and file-ID references in the same request.
Only supported for grok-imagine models.
user: A unique identifier representing your end-user.
image_format: The format of the image to return ("url" or "base64"). Defaults to "url".
aspect_ratio: The aspect ratio of the image to generate.
resolution: The image resolution to generate ("1k" or "2k").
storage_options: Persist the result to the Files API. Accepts a dict
with a required ``filename`` and optional ``expires_after`` and ``public_url`` keys.
Set ``public_url`` to also create a publicly shareable URL.
Examples::

storage_options={"filename": "output.png"} # store privately, no expiry
storage_options={"filename": "output.png", "expires_after": 7200} # auto-delete after 2h
storage_options={"filename": "output.png", "public_url": True} # + shareable URL
storage_options={"filename": "output.png", "public_url": {"expires_after": 86400}} # + expiring URL

Returns:
A `BatchRequest` proto ready to be added to a batch.
Expand Down Expand Up @@ -91,11 +117,14 @@ def prepare(
prompt,
model,
image_url=image_url,
image_file_id=image_file_id,
image_urls=image_urls,
image_file_ids=image_file_ids,
user=user,
image_format=image_format,
aspect_ratio=aspect_ratio,
resolution=resolution,
storage_options=storage_options,
)
return batch_pb2.BatchRequest(
image_request=request,
Expand All @@ -108,23 +137,37 @@ async def sample(
model: Union[ImageGenerationModel, str],
*,
image_url: Optional[str] = None,
image_file_id: Optional[str] = None,
image_urls: Optional[Sequence[str]] = None,
image_file_ids: Optional[Sequence[str]] = None,
user: Optional[str] = None,
image_format: Optional[ImageFormat] = None,
aspect_ratio: Optional[ImageAspectRatio] = None,
resolution: Optional[ImageResolution] = None,
storage_options: Optional[Union[StorageOptions, image_pb2.StorageOptions]] = None,
) -> "ImageResponse":
"""Samples a single image asynchronously based on the provided prompt.

Args:
prompt: The prompt to generate an image from.
model: The model to use for image generation.
image_url: The URL or base64-encoded string of an input image to use as a starting point for generation.
This field cannot be set together with `image_urls`.
This field cannot be set together with `image_file_id`, `image_urls`, or `image_file_ids`.
Only supported for grok-imagine models.
image_file_id: The file ID of an input image to use as a starting point for generation.
Cannot be set together with `image_url`, `image_urls`, or `image_file_ids`.
Only supported for grok-imagine models.
image_urls: Optional list of input images for multi-reference image editing.
Each image is a URL or base64-encoded string, matching the `image_url` format.
This field cannot be set together with `image_url`.
Cannot be set together with single-image params (`image_url` /
`image_file_id`). May be combined with `image_file_ids` to mix
URL/base64 and file-ID references in the same request; file IDs
are appended first in that case.
Only supported for grok-imagine models.
image_file_ids: Optional list of input image file IDs for multi-reference image editing.
Cannot be set together with single-image params (`image_url` /
`image_file_id`). May be combined with `image_urls` to mix
URL/base64 and file-ID references in the same request.
Only supported for grok-imagine models.
user: A unique identifier representing your end-user, which can help xAI to monitor and detect abuse.
image_format: The format of the image to return. One of:
Expand All @@ -150,6 +193,15 @@ async def sample(
- `"1k"`: ~1 megapixel total. Dimensions vary by aspect ratio.
- `"2k"`: ~4 megapixels total. Dimensions vary by aspect ratio.
Only supported for grok-imagine models.
storage_options: Persist the result to the Files API. Accepts a dict
with a required ``filename`` and optional ``expires_after`` and ``public_url`` keys.
Set ``public_url`` to also create a publicly shareable URL.
Examples::

storage_options={"filename": "output.png"} # store privately, no expiry
storage_options={"filename": "output.png", "expires_after": 7200} # auto-delete after 2h
storage_options={"filename": "output.png", "public_url": True} # + shareable URL
storage_options={"filename": "output.png", "public_url": {"expires_after": 86400}} # + expiring URL

Returns:
An `ImageResponse` object allowing access to the generated image.
Expand All @@ -158,11 +210,14 @@ async def sample(
prompt,
model,
image_url=image_url,
image_file_id=image_file_id,
image_urls=image_urls,
image_file_ids=image_file_ids,
user=user,
image_format=image_format,
aspect_ratio=aspect_ratio,
resolution=resolution,
storage_options=storage_options,
)
with tracer.start_as_current_span(
name=f"image.sample {model}",
Expand All @@ -181,11 +236,14 @@ async def sample_batch(
n: int,
*,
image_url: Optional[str] = None,
image_file_id: Optional[str] = None,
image_urls: Optional[Sequence[str]] = None,
image_file_ids: Optional[Sequence[str]] = None,
user: Optional[str] = None,
image_format: Optional[ImageFormat] = None,
aspect_ratio: Optional[ImageAspectRatio] = None,
resolution: Optional[ImageResolution] = None,
storage_options: Optional[Union[StorageOptions, image_pb2.StorageOptions]] = None,
) -> Sequence["ImageResponse"]:
"""Samples a batch of images asynchronously based on the provided prompt.

Expand All @@ -194,11 +252,22 @@ async def sample_batch(
model: The model to use for image generation.
n: The number of images to generate.
image_url: The URL or base64-encoded string of an input image to use as a starting point for generation.
This field cannot be set together with `image_urls`.
This field cannot be set together with `image_file_id`, `image_urls`, or `image_file_ids`.
Only supported for grok-imagine models.
image_file_id: The file ID of an input image to use as a starting point for generation.
Cannot be set together with `image_url`, `image_urls`, or `image_file_ids`.
Only supported for grok-imagine models.
image_urls: Optional list of input images for multi-reference image editing.
Each image is a URL or base64-encoded string, matching the `image_url` format.
This field cannot be set together with `image_url`.
Cannot be set together with single-image params (`image_url` /
`image_file_id`). May be combined with `image_file_ids` to mix
URL/base64 and file-ID references in the same request; file IDs
are appended first in that case.
Only supported for grok-imagine models.
image_file_ids: Optional list of input image file IDs for multi-reference image editing.
Cannot be set together with single-image params (`image_url` /
`image_file_id`). May be combined with `image_urls` to mix
URL/base64 and file-ID references in the same request.
Only supported for grok-imagine models.
user: A unique identifier representing your end-user, which can help xAI to monitor and detect abuse.
image_format: The format of the image to return. One of:
Expand All @@ -224,6 +293,15 @@ async def sample_batch(
- `"1k"`: ~1 megapixel total. Dimensions vary by aspect ratio.
- `"2k"`: ~4 megapixels total. Dimensions vary by aspect ratio.
Only supported for grok-imagine models.
storage_options: Persist the results to the Files API. Accepts a dict
with a required ``filename`` and optional ``expires_after`` and ``public_url`` keys.
Set ``public_url`` to also create a publicly shareable URL.
Examples::

storage_options={"filename": "output.png"} # store privately, no expiry
storage_options={"filename": "output.png", "expires_after": 7200} # auto-delete after 2h
storage_options={"filename": "output.png", "public_url": True} # + shareable URL
storage_options={"filename": "output.png", "public_url": {"expires_after": 86400}} # + expiring URL

Returns:
A sequence of `ImageResponse` objects, one for each image generated.
Expand All @@ -233,11 +311,14 @@ async def sample_batch(
model,
n=n,
image_url=image_url,
image_file_id=image_file_id,
image_urls=image_urls,
image_file_ids=image_file_ids,
user=user,
image_format=image_format,
aspect_ratio=aspect_ratio,
resolution=resolution,
storage_options=storage_options,
)
with tracer.start_as_current_span(
name=f"image.sample_batch {model}",
Expand Down
Loading
Loading