rig_core/providers/gemini/
image_generation.rs

1//! Image generation through Gemini's `generateContent` endpoint. The
2//! request's width and height pick the aspect ratio, its images are
3//! references the model draws from, and [`ImageOptions`] sets the size and
4//! asks for text beside the images. [`parts_of`] reads every image and text
5//! part of a reply.
6//!
7//! ```no_run
8//! use rig_core::image_generation::ImageGenerationRequestBuilder;
9//! use rig_core::providers::gemini::{Gemini, ImageSize};
10//! use rig_core::providers::gemini::image_generation::{GEMINI_3_1_FLASH_IMAGE, ImageOptions};
11//!
12//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
13//! let model = Gemini::from_env()?.image_generation(GEMINI_3_1_FLASH_IMAGE);
14//! let request = ImageGenerationRequestBuilder::new("A banana in a tuxedo")
15//!     .width(16)
16//!     .height(9)
17//!     .additional_params(ImageOptions::new().image_size(ImageSize::TwoK).params())
18//!     .build();
19//! let image = model.call(request).await?.image;
20//! # let _ = image;
21//! # Ok(())
22//! # }
23//! ```
24
25use super::ImageSize;
26use super::completion::usage_of;
27use crate::error::{EncodeError, ProviderError};
28use crate::image_generation;
29use crate::image_generation::ImageGenerationRequest;
30use crate::json_utils::Lenient;
31use crate::message::{DocumentSourceKind, Image};
32use crate::operation::ImageGeneration;
33use crate::wire::{Body, Descriptor, DocumentDecoder, Encoded, Framing, Mode, Wire};
34use base64::Engine;
35use base64::prelude::BASE64_STANDARD;
36use serde_json::{Value, json};
37
38/// `gemini-nano-banana-2.1`, Nano Banana 2.1: the general image model.
39pub const GEMINI_NANO_BANANA_2_1: &str = "gemini-nano-banana-2.1";
40/// `gemini-3.1-flash-image`, Nano Banana 2.
41pub const GEMINI_3_1_FLASH_IMAGE: &str = "gemini-3.1-flash-image";
42/// `gemini-3.1-flash-lite-image`, Nano Banana 2 Lite: 1K images only.
43pub const GEMINI_3_1_FLASH_LITE_IMAGE: &str = "gemini-3.1-flash-lite-image";
44/// `gemini-3-pro-image`, Nano Banana Pro.
45pub const GEMINI_3_PRO_IMAGE: &str = "gemini-3-pro-image";
46/// `gemini-2.5-flash-image`, Nano Banana.
47pub const GEMINI_2_5_FLASH_IMAGE: &str = "gemini-2.5-flash-image";
48
49/// The aspect ratios the API takes, as `(width, height)`.
50const ASPECT_RATIOS: [(u32, u32); 14] = [
51    (1, 1),
52    (1, 4),
53    (4, 1),
54    (1, 8),
55    (8, 1),
56    (2, 3),
57    (3, 2),
58    (3, 4),
59    (4, 3),
60    (4, 5),
61    (5, 4),
62    (9, 16),
63    (16, 9),
64    (21, 9),
65];
66
67/// Gemini's image options, sent as the request's `additional_params`
68/// through [`params`](Self::params).
69#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize)]
70#[serde(rename_all = "camelCase")]
71pub struct ImageOptions {
72    #[serde(skip_serializing_if = "Option::is_none")]
73    aspect_ratio: Option<String>,
74    #[serde(skip_serializing_if = "Option::is_none")]
75    image_size: Option<ImageSize>,
76    #[serde(skip)]
77    text: bool,
78}
79
80impl ImageOptions {
81    /// No option set: the aspect ratio follows the request's width and
82    /// height, the size is the model's default, and the reply is images
83    /// only.
84    pub fn new() -> Self {
85        Self::default()
86    }
87
88    /// `imageConfig.aspectRatio`, such as `"21:9"`, in place of the one the
89    /// width and height pick.
90    pub fn aspect_ratio(mut self, ratio: impl Into<String>) -> Self {
91        self.aspect_ratio = Some(ratio.into());
92        self
93    }
94
95    /// `imageConfig.imageSize`.
96    pub fn image_size(mut self, size: ImageSize) -> Self {
97        self.image_size = Some(size);
98        self
99    }
100
101    /// Ask for text parts beside the images, which [`parts_of`] reads.
102    pub fn with_text(mut self) -> Self {
103        self.text = true;
104        self
105    }
106
107    /// The options as `additional_params`: a `generationConfig` the wire
108    /// merges over its own.
109    pub fn params(self) -> Value {
110        let modalities = self.text.then(|| json!(["TEXT", "IMAGE"]));
111        let mut config = json!({ "imageConfig": self });
112        if let (Some(modalities), Some(config)) = (modalities, config.as_object_mut()) {
113            config.insert("responseModalities".to_owned(), modalities);
114        }
115        json!({ "generationConfig": config })
116    }
117}
118
119/// One part of an image reply.
120#[derive(Clone, Debug, PartialEq, Eq)]
121pub enum ImagePart {
122    /// Text the model wrote.
123    Text(String),
124    /// An image, decoded from base64.
125    Image {
126        /// Its media type, such as `image/png`.
127        mime_type: String,
128        /// Its bytes.
129        data: Vec<u8>,
130    },
131}
132
133/// Every text and image part of the first candidate in `reply`, in reply
134/// order. Thoughts, interim images included, and images without data are
135/// left out.
136///
137/// # Errors
138///
139/// When an image's data is not base64.
140pub fn parts_of(reply: &Value) -> Result<Vec<ImagePart>, ProviderError> {
141    let parts = reply
142        .arr("candidates")
143        .first()
144        .and_then(|candidate| candidate.get("content"))
145        .map(|content| content.arr("parts"))
146        .unwrap_or_default();
147    let mut out = Vec::new();
148    for part in parts
149        .iter()
150        .filter(|part| part.bool("thought") != Some(true))
151    {
152        if let Some(text) = part.str("text") {
153            out.push(ImagePart::Text(text.to_owned()));
154        }
155        let Some(blob) = part.get("inlineData") else {
156            continue;
157        };
158        let mime_type = blob.str("mimeType").unwrap_or_default();
159        let Some(data) = blob.str("data").filter(|_| mime_type.starts_with("image/")) else {
160            continue;
161        };
162        let data = BASE64_STANDARD.decode(data).map_err(|err| {
163            ProviderError::Response(format!("Gemini image data was not valid base64: {err}"))
164        })?;
165        out.push(ImagePart::Image {
166            mime_type: mime_type.to_owned(),
167            data,
168        });
169    }
170    Ok(out)
171}
172
173/// The first image in `reply`, with its usage and identity.
174///
175/// # Errors
176///
177/// When `reply` holds no image, or data that is not base64.
178pub fn image_of(reply: &Value) -> Result<image_generation::ImageGenerationResponse, ProviderError> {
179    let image = parts_of(reply)?
180        .into_iter()
181        .find_map(|part| match part {
182            ImagePart::Image { data, .. } => Some(data),
183            ImagePart::Text(_) => None,
184        })
185        .ok_or_else(|| {
186            ProviderError::Response(
187                "Gemini image generation response did not include image data".into(),
188            )
189        })?;
190    Ok(image_generation::ImageGenerationResponse {
191        model: reply.str("modelVersion").map(str::to_owned),
192        response_id: Some(reply.str("responseId").unwrap_or_default().to_owned()),
193        usage: reply.get("usageMetadata").map(usage_of).unwrap_or_default(),
194        ..image_generation::ImageGenerationResponse::new(image)
195    })
196}
197
198fn create_request_body(request: ImageGenerationRequest) -> Result<Value, EncodeError> {
199    let mut parts = vec![json!({ "text": request.prompt })];
200    for image in request.images {
201        parts.push(reference(image)?);
202    }
203    let mut image_config = serde_json::Map::new();
204    if let Some(ratio) = aspect_ratio(request.width, request.height) {
205        image_config.insert("aspectRatio".to_owned(), json!(ratio));
206    }
207    let mut body = json!({
208        "contents": [{ "role": "user", "parts": parts }],
209        "toolConfig": null,
210        "generationConfig": { "responseModalities": ["IMAGE"], "imageConfig": image_config },
211        "safetySettings": null,
212        "systemInstruction": null,
213    });
214    match (&mut body, request.additional_params) {
215        (Value::Object(body), Some(Value::Object(params))) => {
216            crate::completion::options::deep_merge(body, params);
217        }
218        (body, Some(params)) => *body = params,
219        (_, None) => {}
220    }
221    Ok(body)
222}
223
224/// A reference image as a part: a URL by reference, and data inline.
225fn reference(image: Image) -> Result<Value, EncodeError> {
226    let source = match image.data {
227        DocumentSourceKind::Raw(bytes) => DocumentSourceKind::Base64(BASE64_STANDARD.encode(bytes)),
228        source => source,
229    };
230    super::completion::media(super::completion::mime(image.media_type), &source, false)
231}
232
233/// The aspect ratio the API takes that `width` by `height` has; `None` for
234/// a zero side or a ratio it does not take, which leaves the choice to the
235/// model.
236fn aspect_ratio(width: u32, height: u32) -> Option<String> {
237    let (width, height) = (u64::from(width), u64::from(height));
238    ASPECT_RATIOS
239        .iter()
240        .find(|&&(w, h)| width > 0 && width * u64::from(h) == height * u64::from(w))
241        .map(|(w, h)| format!("{w}:{h}"))
242}
243
244/// The image generation wire: `POST /v1beta/models/{model}:generateContent`.
245///
246/// Both [`Mode`]s request image output and decode a whole response document.
247#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
248pub struct Images {
249    /// The provider this wire speaks to.
250    pub provider: super::GeminiConfig,
251    /// Name of the model, for example [`GEMINI_3_1_FLASH_IMAGE`].
252    pub model: String,
253}
254
255impl Images {
256    /// The image generation wire for `model`.
257    pub fn new(provider: super::GeminiConfig, model: impl Into<String>) -> Self {
258        Self {
259            provider,
260            model: model.into(),
261        }
262    }
263}
264
265impl Wire for Images {
266    type Op = ImageGeneration;
267    type Payload = crate::wire::Encoded;
268    type Frame = crate::wire::WireFrame;
269    type Event = Value;
270    type Decoder<'id> = DocumentDecoder<ImageGeneration, Value>;
271    type Reassembler = crate::wire::document::Unreassembled;
272
273    fn describe(&self) -> Descriptor<'_> {
274        Descriptor::new(super::PROVIDER_NAME).model(self.model.as_str())
275    }
276
277    fn encode(&self, request: ImageGenerationRequest, _mode: Mode) -> Result<Encoded, EncodeError> {
278        let body = serde_json::to_vec(&create_request_body(request)?)?;
279        let request = http::Request::post(format!(
280            "{}/v1beta/models/{}:generateContent?key={}",
281            self.provider.base_url,
282            self.model,
283            self.provider.api_key.expose()
284        ))
285        .header(http::header::CONTENT_TYPE, "application/json")
286        .body(Body::Bytes(body))?;
287        // Gemini reports no transport request-id header.
288        Ok(Encoded::new(request, Framing::Whole))
289    }
290
291    /// The first non-thought image; missing image data and invalid base64
292    /// are response errors.
293    fn decoder<'id>(&self) -> Self::Decoder<'id> {
294        DocumentDecoder::new(super::completion::classify_reply, |reply| image_of(&reply))
295    }
296}
297
298impl super::GeminiConfig {
299    /// The image generation wire.
300    pub(crate) fn image_generation(&self, model: impl Into<String>) -> Images {
301        Images::new(self.clone(), model)
302    }
303}
304
305#[cfg(test)]
306mod tests;
rig_core/providers/gemini/image_generation.rs

rig_core/providers/gemini/
image_generation.rs
