feat: impl imagen (#16)

Author: wiseaidev

README.md

@@ -100,6 +100,12 @@ gems info
 gems list
 ```
 
+### Generate an Image
+
+```sh
+gems imagen -t "Hi, can you create a 3d rendered image of a pig with wings and a top hat flying over a happy futuristic scifi city with lots of greenery?"
+```
+
 ## 🎨 Options
 
 | Option                   | Description                                              |
@@ -115,6 +121,7 @@ gems list
 | `generate`              | Generate creative content.                               |
 | `vision`                | Analyze an image and generate content from text.         |
 | `stream`                | Stream the generation of content.                        |
+| `imagen`                | Generate an image.                                       |
 | `count`                 | Count the number of tokens in a text.                    |
 | `embed`                 | Embed content into a specified model.                    |
 | `batch`                 | Batch embed multiple contents.                           |

src/chat.rs

@@ -32,6 +32,7 @@ impl Chats {
         let request_body = GeminiRequest {
             model: params.model.to_string(),
             contents: vec![content],
+            config: None,
         };
 
         let req = self

src/cli.rs

@@ -101,6 +101,7 @@ pub enum Command {
     Batch(Batch),
     Info(Info),
     List(List),
+    Imagen(Imagen),
 }
 
 #[cfg(feature = "cli")]
@@ -161,3 +162,11 @@ pub struct Vision {
     #[arg(short, long, default_value_t = String::from("What is this picture?"))]
     pub text: String,
 }
+
+#[cfg(feature = "cli")]
+#[derive(Args, Debug, Clone)]
+pub struct Imagen {
+    /// The text to generate image from.
+    #[arg(short, long, default_value_t = String::from("Hi, step bro... I need help generating a happy, humble, bumble Rustacean. he's stuck in the shower and won't compile."))]
+    pub text: String,
+}

src/client.rs

@@ -1,5 +1,6 @@
 use crate::chat::Chats;
 use crate::embed::Embeddings;
+use crate::imagen::Images;
 use crate::models::Model;
 use crate::models::Models;
 use crate::stream::Streaming;
@@ -108,6 +109,12 @@ impl CTrait for Client {
             client: self.clone(),
         }
     }
+
+    fn images(&self) -> Images {
+        Images {
+            client: self.clone(),
+        }
+    }
 }
 
 #[derive(Default)]

src/imagen.rs

@@ -0,0 +1,57 @@
+use crate::client::Client;
+use crate::messages::Message;
+use crate::models::Model;
+use crate::requests::GenerationConfig;
+use crate::requests::{Content, GeminiRequest};
+use crate::responses::ImagenResponse;
+use crate::traits::CTrait;
+use crate::utils::extract_image_or_text;
+use anyhow::{anyhow, Result};
+use derive_builder::Builder;
+use reqwest::Method;
+
+#[derive(Clone)]
+pub struct Images {
+    pub client: Client,
+}
+
+#[derive(Builder, Clone)]
+#[builder(setter(into))]
+pub struct ImageGen {
+    pub model: Model,
+    pub input: Message,
+}
+
+impl Images {
+    pub async fn generate(&self, params: ImageGen) -> Result<Vec<u8>> {
+        let content = Content {
+            parts: vec![params.input.to_part()],
+        };
+        let request_body = GeminiRequest {
+            model: params.model.to_string(),
+            contents: vec![content],
+            config: Some(GenerationConfig {
+                response_modalities: vec!["Text".into(), "Image".into()],
+            }),
+        };
+
+        let req = self
+            .client
+            .request(Method::POST, "generateContent")?
+            .json(&request_body);
+
+        let res = req.send().await?;
+        let json: ImagenResponse = res.json().await?;
+
+        let parts = json
+            .candidates
+            .ok_or_else(|| anyhow!("Missing candidates"))?
+            .first()
+            .ok_or_else(|| anyhow!("No candidate response"))?
+            .content
+            .parts
+            .clone();
+
+        extract_image_or_text(&parts)
+    }
+}

src/lib.rs

@@ -4,6 +4,7 @@
 pub mod chat;
 pub mod client;
 pub mod embed;
+pub mod imagen;
 pub mod messages;
 pub mod models;
 pub mod requests;

src/main.rs

@@ -14,6 +14,7 @@ async fn main() -> Result<()> {
         use gems::cli::{Cli, Command};
         use gems::embed::BatchEmbeddingBuilder;
         use gems::embed::EmbeddingBuilder;
+        use gems::imagen::ImageGenBuilder;
         use gems::messages::Content;
         use gems::messages::Message;
         use gems::models::ModBuilder;
@@ -167,6 +168,22 @@ async fn main() -> Result<()> {
                 let models = gemini_client.models().list().await?;
                 models.print();
             }
+            Command::Imagen(cmd) => {
+                gemini_client.set_model(Model::FlashExpImage);
+
+                let params = ImageGenBuilder::default()
+                    .input(Message::User {
+                        content: Content::Text(cmd.text),
+                        name: None,
+                    })
+                    .model(Model::FlashExpImage)
+                    .build()
+                    .unwrap();
+
+                let image_data = gemini_client.images().generate(params).await?;
+
+                std::fs::write("output.png", &image_data)?;
+            }
         }
     }
     Ok(())

src/models.rs

@@ -20,6 +20,7 @@ pub enum Model {
     Imagen3,
     Veo2,
     Flash20Live,
+    FlashExpImage,
 }
 
 #[allow(clippy::to_string_trait_impl)]
@@ -36,6 +37,7 @@ impl ToString for Model {
             Model::Imagen3 => "imagen-3.0-generate-002",
             Model::Veo2 => "veo-2.0-generate-001",
             Model::Flash20Live => "gemini-2.0-flash-live-001",
+            Model::FlashExpImage => "gemini-2.0-flash-exp-image-generation",
         }
         .to_string()
     }
@@ -56,6 +58,7 @@ impl FromStr for Model {
             "imagen-3.0-generate-002" => Ok(Model::Imagen3),
             "veo-2.0-generate-001" => Ok(Model::Veo2),
             "gemini-2.0-flash-live-001" => Ok(Model::Flash20Live),
+            "gemini-2.0-flash-exp-image-generation" => Ok(Model::Flash20Live),
             _ => Err(anyhow!("Unknown model: {}", s)),
         }
     }

src/requests.rs

@@ -8,7 +8,11 @@ pub struct GeminiRequest {
 
     /// List of content items for generation.
     pub contents: Vec<Content>,
+
+    #[serde(rename = "generationConfig", skip_serializing_if = "Option::is_none")]
+    pub config: Option<GenerationConfig>,
 }
+
 /// Request structure for content embedding.
 #[derive(Debug, Serialize, Deserialize)]
 pub struct GeminiEmbedRequest {
@@ -32,7 +36,7 @@ pub struct Content {
 }
 
 /// Define an enum to represent different types of parts in the content.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 #[serde(untagged)]
 pub enum Part {
     /// Represents a text part in the content.
@@ -64,10 +68,16 @@ pub struct Candidate {
 }
 
 /// Structure representing the image part of the Gemini request.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct ImageContent {
     /// The MIME type of the image.
     pub mime_type: String,
     /// The actual image data in a base64-encoded string.
     pub data: String,
 }
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct GenerationConfig {
+    #[serde(rename = "responseModalities")]
+    pub response_modalities: Vec<String>,
+}

src/responses.rs

@@ -1,4 +1,4 @@
-use crate::requests::Candidate;
+use crate::requests::Candidate as ReqCandidate;
 use serde::{Deserialize, Serialize};
 
 /// Response structure for content embedding.
@@ -120,5 +120,62 @@ impl ModelsResponse {
 #[derive(Debug, Deserialize)]
 pub struct GeminiResponse {
     /// List of generated candidates.
+    pub candidates: Option<Vec<ReqCandidate>>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ImagenResponse {
     pub candidates: Option<Vec<Candidate>>,
+    pub usage_metadata: Option<UsageMetadata>,
+    pub model_version: Option<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct Candidate {
+    pub content: Content,
+    pub finish_reason: Option<String>,
+    pub index: Option<i32>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct Content {
+    pub parts: Vec<Part>,
+    pub role: Option<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+#[serde(untagged)]
+pub enum Part {
+    Text {
+        text: String,
+    },
+    Image {
+        #[serde(rename = "inlineData")]
+        inline_data: ImageContent,
+    },
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+#[serde(rename_all = "camelCase")]
+pub struct ImageContent {
+    pub mime_type: String,
+    pub data: String,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct UsageMetadata {
+    pub prompt_token_count: Option<i32>,
+    pub total_token_count: Option<i32>,
+    pub prompt_tokens_details: Option<Vec<PromptTokenDetail>>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct PromptTokenDetail {
+    pub modality: Option<String>,
+    pub token_count: Option<i32>,
 }

src/stream.rs

@@ -28,6 +28,7 @@ impl Streaming {
             contents: vec![Content {
                 parts: vec![params.input.to_part()],
             }],
+            config: None,
         };
 
         let req = self

src/tokens.rs

@@ -28,6 +28,7 @@ impl Tokens {
             contents: vec![Content {
                 parts: vec![params.input.to_part()],
             }],
+            config: None,
         };
 
         let req = self

src/traits.rs

@@ -1,5 +1,6 @@
 use crate::chat::Chats;
 use crate::embed::Embeddings;
+use crate::imagen::Images;
 use crate::models::Model;
 use crate::models::Models;
 use crate::stream::Streaming;
@@ -21,4 +22,5 @@ pub trait CTrait {
     fn vision(&self) -> Visions;
     fn stream(&self) -> Streaming;
     fn models(&self) -> Models;
+    fn images(&self) -> Images;
 }

src/utils.rs

@@ -1,4 +1,5 @@
-use anyhow::Result;
+use crate::responses::Part;
+use anyhow::{anyhow, Result};
 use base64::{engine::general_purpose::STANDARD, Engine as _};
 use std::fs::File;
 use std::io::Read;
@@ -97,3 +98,21 @@ pub fn load_and_encode_image(file_path: &str) -> Result<String> {
     let base64_encoded = STANDARD.encode(&buffer);
     Ok(base64_encoded)
 }
+
+pub fn extract_image_or_text(parts: &[Part]) -> Result<Vec<u8>> {
+    if let Some(base64_data) = parts.iter().find_map(|part| match part {
+        Part::Image { inline_data } => Some(inline_data.data.clone()),
+        _ => None,
+    }) {
+        let image_bytes = STANDARD.decode(&base64_data)?;
+        Ok(image_bytes)
+    }
+    else if let Some(text) = parts.iter().find_map(|part| match part {
+        Part::Text { text } => Some(text.clone()),
+        _ => None,
+    }) {
+        Err(anyhow!("Expected image but got only text: {}", text))
+    } else {
+        Err(anyhow!("No image or text found in response"))
+    }
+}

src/vision.rs

@@ -47,6 +47,7 @@ impl Visions {
             contents: vec![crate::requests::Content {
                 parts: vec![input_part, image_part],
             }],
+            config: None,
         };
 
         let req = self