aither_gemini/
image.rs

1use aither_core::image::{Data as ImageData, ImageGenerator, Prompt, Size};
2use futures_core::Stream;
3use futures_lite::StreamExt;
4
5use crate::{
6    client::call_generate,
7    config::Gemini,
8    error::GeminiError,
9    types::{GeminiContent, GenerateContentRequest, GenerationConfig, ImageConfig, Part},
10};
11
12impl ImageGenerator for Gemini {
13    type Error = GeminiError;
14
15    fn create(
16        &self,
17        prompt: Prompt,
18        size: Size,
19    ) -> impl Stream<Item = Result<ImageData, Self::Error>> + Send {
20        let cfg = self.config();
21        let text = prompt.text().to_owned();
22        let images = prompt.images().to_vec();
23        let aspect_ratio = map_size_to_aspect_ratio(size);
24
25        futures_lite::stream::iter(vec![async move {
26            let model = cfg.image_model.clone().ok_or_else(|| {
27                GeminiError::Api("image generation is disabled for this Gemini backend".into())
28            })?;
29            let mut parts = vec![Part::text(text)];
30            for image in images {
31                parts.push(Part::inline_image(image));
32            }
33
34            let request = GenerateContentRequest {
35                system_instruction: None,
36                contents: vec![GeminiContent::with_parts("user", parts)],
37                generation_config: Some(GenerationConfig {
38                    image_config: Some(ImageConfig {
39                        aspect_ratio: Some(aspect_ratio),
40                        image_size: None,
41                    }),
42                    ..GenerationConfig::default()
43                }),
44                tools: Vec::new(),
45                tool_config: None,
46                safety_settings: Vec::new(),
47                cached_content: None,
48            };
49            call_generate(cfg, &model, request).await
50        }])
51        .then(|fut| fut)
52        .filter_map(Result::ok)
53        .flat_map(|response| {
54            let parts = response
55                .primary_candidate()
56                .and_then(|candidate| candidate.content.as_ref())
57                .map(|content| content.parts.clone())
58                .unwrap_or_default();
59            futures_lite::stream::iter(parts)
60                .filter_map(|part| {
61                    part.inline_data
62                        .as_ref()
63                        .and_then(|inline| inline.decode().ok())
64                })
65                .map(Ok)
66        })
67    }
68
69    fn edit(
70        &self,
71        prompt: Prompt,
72        mask: &[u8],
73    ) -> impl Stream<Item = Result<ImageData, Self::Error>> + Send {
74        let cfg = self.config();
75        let mask_bytes = mask.to_vec();
76        let text = prompt.text().to_owned();
77        let base_image = prompt.images().first().cloned();
78
79        futures_lite::stream::iter(vec![async move {
80            let model = cfg.image_model.clone().ok_or_else(|| {
81                GeminiError::Api("image generation is disabled for this Gemini backend".into())
82            })?;
83            let base_image = base_image.ok_or_else(|| {
84                GeminiError::Api(
85                    "image editing requires Prompt::with_image to supply a base image".into(),
86                )
87            })?;
88            let mut parts = vec![Part::inline_image(base_image)];
89            if !mask_bytes.is_empty() {
90                parts.push(Part::inline_mask(mask_bytes));
91            }
92            parts.push(Part::text(format!(
93                "Apply the mask and follow these instructions: {text}"
94            )));
95            let request = GenerateContentRequest {
96                system_instruction: None,
97                contents: vec![GeminiContent::with_parts("user", parts)],
98                generation_config: Some(GenerationConfig::default()),
99                tools: Vec::new(),
100                tool_config: None,
101                safety_settings: Vec::new(),
102                cached_content: None,
103            };
104            call_generate(cfg, &model, request).await
105        }])
106        .then(|fut| fut)
107        .filter_map(Result::ok)
108        .flat_map(|response| {
109            let parts = response
110                .primary_candidate()
111                .and_then(|candidate| candidate.content.as_ref())
112                .map(|content| content.parts.clone())
113                .unwrap_or_default();
114            futures_lite::stream::iter(parts)
115                .filter_map(|part| {
116                    part.inline_data
117                        .as_ref()
118                        .and_then(|inline| inline.decode().ok())
119                })
120                .map(Ok)
121        })
122    }
123}
124
125fn map_size_to_aspect_ratio(size: Size) -> String {
126    if size.width() == size.height() {
127        "1:1".into()
128    } else if size.width() > size.height() {
129        if size.width() * 9 == size.height() * 16 {
130            "16:9".into()
131        } else if size.width() * 3 == size.height() * 4 {
132            "4:3".into()
133        } else {
134            "16:9".into() // Default landscape
135        }
136    } else if size.width() * 16 == size.height() * 9 {
137        "9:16".into()
138    } else if size.width() * 4 == size.height() * 3 {
139        "3:4".into()
140    } else {
141        "9:16".into() // Default portrait
142    }
143}
aither_gemini/image.rs

aither_gemini/
image.rs
