Skip to main content

edgequake_llm/imagegen/
types.rs

1//! Shared types for image generation.
2
3use std::collections::HashMap;
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value as JsonValue;
7
8/// Named aspect ratio presets shared across providers.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
10#[serde(rename_all = "snake_case")]
11pub enum AspectRatio {
12    #[default]
13    Auto,
14    Square,
15    SquareHd,
16    Landscape43,
17    Landscape169,
18    Ultrawide,
19    Portrait43,
20    Portrait169,
21    Frame54,
22    Frame45,
23    Print32,
24    Print23,
25    Extreme41,
26    Extreme14,
27    Extreme81,
28    Extreme18,
29}
30
31impl AspectRatio {
32    /// Default pixel size used when the provider does not return explicit dimensions.
33    pub fn default_dimensions(self) -> (u32, u32) {
34        match self {
35            Self::Auto => (1024, 1024),
36            Self::Square => (1024, 1024),
37            Self::SquareHd => (2048, 2048),
38            Self::Landscape43 => (1024, 768),
39            Self::Landscape169 => (1280, 720),
40            Self::Ultrawide => (2560, 1080),
41            Self::Portrait43 => (768, 1024),
42            Self::Portrait169 => (720, 1280),
43            Self::Frame54 => (1280, 1024),
44            Self::Frame45 => (1024, 1280),
45            Self::Print32 => (1152, 768),
46            Self::Print23 => (768, 1152),
47            Self::Extreme41 => (2048, 512),
48            Self::Extreme14 => (512, 2048),
49            Self::Extreme81 => (3072, 384),
50            Self::Extreme18 => (384, 3072),
51        }
52    }
53
54    /// Vertex Imagen ratio name.
55    pub fn as_vertex_str(self) -> &'static str {
56        match self {
57            Self::Auto | Self::Square | Self::SquareHd => "1:1",
58            Self::Landscape43 => "4:3",
59            Self::Landscape169 | Self::Ultrawide => "16:9",
60            Self::Portrait43 => "3:4",
61            Self::Portrait169 => "9:16",
62            Self::Frame54 => "5:4",
63            Self::Frame45 => "4:5",
64            Self::Print32 => "3:2",
65            Self::Print23 => "2:3",
66            Self::Extreme41 => "4:1",
67            Self::Extreme14 => "1:4",
68            Self::Extreme81 => "8:1",
69            Self::Extreme18 => "1:8",
70        }
71    }
72
73    /// FAL image size preset.
74    pub fn as_fal_str(self) -> &'static str {
75        match self {
76            Self::Auto | Self::Square => "square",
77            Self::SquareHd => "square_hd",
78            Self::Landscape43 | Self::Frame54 | Self::Print32 => "landscape_4_3",
79            Self::Landscape169 | Self::Ultrawide | Self::Extreme41 | Self::Extreme81 => {
80                "landscape_16_9"
81            }
82            Self::Portrait43 | Self::Frame45 | Self::Print23 => "portrait_4_3",
83            Self::Portrait169 | Self::Extreme14 | Self::Extreme18 => "portrait_16_9",
84        }
85    }
86
87    /// Gemini image ratio name.
88    pub fn as_gemini_str(self) -> &'static str {
89        match self {
90            Self::Auto => "auto",
91            Self::Square | Self::SquareHd => "1:1",
92            Self::Landscape43 => "4:3",
93            Self::Landscape169 => "16:9",
94            Self::Ultrawide => "21:9",
95            Self::Portrait43 => "3:4",
96            Self::Portrait169 => "9:16",
97            Self::Frame54 => "5:4",
98            Self::Frame45 => "4:5",
99            Self::Print32 => "3:2",
100            Self::Print23 => "2:3",
101            Self::Extreme41 => "4:1",
102            Self::Extreme14 => "1:4",
103            Self::Extreme81 => "8:1",
104            Self::Extreme18 => "1:8",
105        }
106    }
107}
108
109/// Output image format.
110#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
111#[serde(rename_all = "lowercase")]
112pub enum ImageFormat {
113    #[default]
114    Jpeg,
115    Png,
116    Webp,
117}
118
119impl ImageFormat {
120    pub fn mime_type(self) -> &'static str {
121        match self {
122            Self::Jpeg => "image/jpeg",
123            Self::Png => "image/png",
124            Self::Webp => "image/webp",
125        }
126    }
127
128    pub fn extension(self) -> &'static str {
129        match self {
130            Self::Jpeg => "jpg",
131            Self::Png => "png",
132            Self::Webp => "webp",
133        }
134    }
135
136    pub fn as_fal_str(self) -> &'static str {
137        match self {
138            Self::Jpeg => "jpeg",
139            Self::Png => "png",
140            Self::Webp => "webp",
141        }
142    }
143}
144
145/// Provider-agnostic content safety level.
146#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
147#[serde(rename_all = "snake_case")]
148pub enum SafetyLevel {
149    BlockNone,
150    BlockLow,
151    #[default]
152    BlockMedium,
153    BlockHigh,
154}
155
156/// Gemini image size.
157#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
158pub enum ImageResolution {
159    Half,
160    #[default]
161    OneK,
162    TwoK,
163    FourK,
164}
165
166impl ImageResolution {
167    pub fn as_gemini_str(self) -> &'static str {
168        match self {
169            Self::Half => "512",
170            Self::OneK => "1K",
171            Self::TwoK => "2K",
172            Self::FourK => "4K",
173        }
174    }
175
176    pub fn as_nano_banana_str(self) -> &'static str {
177        match self {
178            Self::Half => "0.5K",
179            Self::OneK => "1K",
180            Self::TwoK => "2K",
181            Self::FourK => "4K",
182        }
183    }
184}
185
186/// Gemini reasoning depth.
187#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
188#[serde(rename_all = "snake_case")]
189pub enum ThinkingLevel {
190    #[default]
191    Minimal,
192    High,
193}
194
195impl ThinkingLevel {
196    pub fn as_gemini_str(self) -> &'static str {
197        match self {
198            Self::Minimal => "minimal",
199            Self::High => "high",
200        }
201    }
202
203    pub fn as_gemini_api_label(self) -> &'static str {
204        match self {
205            Self::Minimal => "Minimal",
206            Self::High => "High",
207        }
208    }
209}
210
211/// Shared provider options.
212#[derive(Debug, Clone, Default, Serialize, Deserialize)]
213pub struct ImageGenOptions {
214    pub count: Option<u8>,
215    pub aspect_ratio: Option<AspectRatio>,
216    pub width: Option<u32>,
217    pub height: Option<u32>,
218    pub seed: Option<u64>,
219    pub negative_prompt: Option<String>,
220    pub guidance_scale: Option<f32>,
221    pub output_format: Option<ImageFormat>,
222    pub safety_level: Option<SafetyLevel>,
223    pub enhance_prompt: Option<bool>,
224    pub resolution: Option<ImageResolution>,
225    pub enable_web_search: Option<bool>,
226    pub thinking_level: Option<ThinkingLevel>,
227    #[serde(default)]
228    pub reference_images: Vec<String>,
229    #[serde(default)]
230    pub extra: HashMap<String, JsonValue>,
231}
232
233impl ImageGenOptions {
234    pub fn aspect_ratio_or_default(&self) -> AspectRatio {
235        self.aspect_ratio.unwrap_or_default()
236    }
237
238    pub fn count_or_default(&self) -> u8 {
239        self.count.unwrap_or(1)
240    }
241
242    pub fn output_format_or_default(&self) -> ImageFormat {
243        self.output_format.unwrap_or_default()
244    }
245
246    pub fn safety_level_or_default(&self) -> SafetyLevel {
247        self.safety_level.unwrap_or_default()
248    }
249
250    pub fn resolution_or_default(&self) -> ImageResolution {
251        self.resolution.unwrap_or_default()
252    }
253}
254
255/// Image generation request.
256#[derive(Debug, Clone, Serialize, Deserialize)]
257pub struct ImageGenRequest {
258    pub prompt: String,
259    #[serde(skip_serializing_if = "Option::is_none")]
260    pub model: Option<String>,
261    #[serde(default)]
262    pub options: ImageGenOptions,
263}
264
265impl ImageGenRequest {
266    pub fn new(prompt: impl Into<String>) -> Self {
267        Self {
268            prompt: prompt.into(),
269            model: None,
270            options: ImageGenOptions::default(),
271        }
272    }
273
274    pub fn with_model(mut self, model: impl Into<String>) -> Self {
275        self.model = Some(model.into());
276        self
277    }
278
279    pub fn with_options(mut self, options: ImageGenOptions) -> Self {
280        self.options = options;
281        self
282    }
283}
284
285/// Generated image payload.
286#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
287pub enum ImageGenData {
288    Bytes(Vec<u8>),
289    Url(String),
290}
291
292/// Single generated image.
293#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
294pub struct GeneratedImage {
295    pub data: ImageGenData,
296    pub width: u32,
297    pub height: u32,
298    pub mime_type: String,
299    pub seed: Option<u64>,
300}
301
302/// Image generation response.
303#[derive(Debug, Clone, Serialize, Deserialize)]
304pub struct ImageGenResponse {
305    pub images: Vec<GeneratedImage>,
306    pub provider: String,
307    pub model: String,
308    pub latency_ms: u64,
309    #[serde(skip_serializing_if = "Option::is_none")]
310    pub enhanced_prompt: Option<String>,
311}