1use std::future::Future;
4use std::str::FromStr;
5
6use bytes::Bytes;
7use serde::Deserialize;
8use serde::Serialize;
9use tokio_util::sync::CancellationToken;
10
11use crate::error::ProviderError;
12use crate::language_model::ResponseMetadata;
13use crate::shared::FileData;
14use crate::shared::Headers;
15use crate::shared::MediaType;
16use crate::shared::ModelId;
17use crate::shared::ProviderId;
18use crate::shared::ProviderMetadata;
19use crate::shared::ProviderOptions;
20use crate::shared::Warning;
21use crate::shared::base64_bytes;
22
23pub trait ImageModel: Send + Sync + 'static {
25 fn provider(&self) -> &ProviderId;
27
28 fn model_id(&self) -> &ModelId;
30
31 fn max_images_per_call(&self) -> Option<usize>;
33
34 fn do_generate(
36 &self,
37 options: ImageOptions,
38 ) -> impl Future<Output = Result<ImageResult, ProviderError>> + Send;
39}
40
41#[derive(Debug, Clone)]
43pub struct ImageOptions {
44 pub prompt: Option<String>,
46 pub n: u32,
48 pub size: Option<ImageSize>,
50 pub aspect_ratio: Option<AspectRatio>,
52 pub seed: Option<u64>,
54 pub files: Vec<ImageFile>,
56 pub mask: Option<ImageFile>,
58 pub provider_options: ProviderOptions,
60 pub headers: Headers,
62 pub cancellation: CancellationToken,
64}
65
66impl ImageOptions {
67 #[must_use]
69 pub fn new(prompt: impl Into<String>) -> Self {
70 Self {
71 prompt: Some(prompt.into()),
72 ..Self::default()
73 }
74 }
75}
76
77impl Default for ImageOptions {
78 fn default() -> Self {
79 Self {
80 prompt: None,
81 n: 1,
82 size: None,
83 aspect_ratio: None,
84 seed: None,
85 files: Vec::new(),
86 mask: None,
87 provider_options: ProviderOptions::new(),
88 headers: Headers::new(),
89 cancellation: CancellationToken::new(),
90 }
91 }
92}
93
94#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
96pub struct ImageFile {
97 pub data: FileData,
99 #[serde(default, skip_serializing_if = "Option::is_none")]
101 pub media_type: Option<MediaType>,
102 #[serde(default, skip_serializing_if = "Option::is_none")]
104 pub provider_options: Option<ProviderOptions>,
105}
106
107#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
109pub struct GeneratedImage {
110 #[serde(with = "base64_bytes")]
112 pub data: Bytes,
113 #[serde(default, skip_serializing_if = "Option::is_none")]
115 pub media_type: Option<MediaType>,
116}
117
118#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
120pub struct ImageUsage {
121 #[serde(default, skip_serializing_if = "Option::is_none")]
123 pub input_tokens: Option<u64>,
124 #[serde(default, skip_serializing_if = "Option::is_none")]
126 pub output_tokens: Option<u64>,
127 #[serde(default, skip_serializing_if = "Option::is_none")]
129 pub total_tokens: Option<u64>,
130}
131
132#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
134pub struct ImageResult {
135 pub images: Vec<GeneratedImage>,
137 #[serde(default, skip_serializing_if = "Option::is_none")]
139 pub is_retryable: Option<bool>,
140 #[serde(default)]
142 pub warnings: Vec<Warning>,
143 #[serde(default, skip_serializing_if = "Option::is_none")]
145 pub provider_metadata: Option<ProviderMetadata>,
146 #[serde(default)]
148 pub response: ResponseMetadata,
149 #[serde(default, skip_serializing_if = "Option::is_none")]
151 pub usage: Option<ImageUsage>,
152}
153
154#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
156#[serde(try_from = "String", into = "String")]
157pub struct ImageSize {
158 pub width: u32,
160 pub height: u32,
162}
163
164impl ImageSize {
165 #[must_use]
167 pub fn new(width: u32, height: u32) -> Self {
168 Self { width, height }
169 }
170}
171
172impl std::fmt::Display for ImageSize {
173 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
174 write!(f, "{}x{}", self.width, self.height)
175 }
176}
177
178impl FromStr for ImageSize {
179 type Err = InvalidDimension;
180
181 fn from_str(text: &str) -> Result<Self, Self::Err> {
182 parse_pair(text, 'x')
183 .map(|(width, height)| Self { width, height })
184 .ok_or_else(|| InvalidDimension {
185 text: text.to_owned(),
186 expected: "WIDTHxHEIGHT",
187 })
188 }
189}
190
191impl TryFrom<String> for ImageSize {
192 type Error = InvalidDimension;
193
194 fn try_from(text: String) -> Result<Self, Self::Error> {
195 text.parse()
196 }
197}
198
199impl From<ImageSize> for String {
200 fn from(size: ImageSize) -> Self {
201 size.to_string()
202 }
203}
204
205#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
207#[serde(try_from = "String", into = "String")]
208pub struct AspectRatio {
209 pub width: u32,
211 pub height: u32,
213}
214
215impl AspectRatio {
216 #[must_use]
218 pub fn new(width: u32, height: u32) -> Self {
219 Self { width, height }
220 }
221}
222
223impl std::fmt::Display for AspectRatio {
224 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
225 write!(f, "{}:{}", self.width, self.height)
226 }
227}
228
229impl FromStr for AspectRatio {
230 type Err = InvalidDimension;
231
232 fn from_str(text: &str) -> Result<Self, Self::Err> {
233 parse_pair(text, ':')
234 .map(|(width, height)| Self { width, height })
235 .ok_or_else(|| InvalidDimension {
236 text: text.to_owned(),
237 expected: "WIDTH:HEIGHT",
238 })
239 }
240}
241
242impl TryFrom<String> for AspectRatio {
243 type Error = InvalidDimension;
244
245 fn try_from(text: String) -> Result<Self, Self::Error> {
246 text.parse()
247 }
248}
249
250impl From<AspectRatio> for String {
251 fn from(ratio: AspectRatio) -> Self {
252 ratio.to_string()
253 }
254}
255
256fn parse_pair(text: &str, separator: char) -> Option<(u32, u32)> {
257 let (left, right) = text.trim().split_once(separator)?;
258 let left: u32 = left.trim().parse().ok()?;
259 let right: u32 = right.trim().parse().ok()?;
260 (left > 0 && right > 0).then_some((left, right))
261}
262
263#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
265#[error("invalid dimension `{text}`: expected `{expected}`")]
266pub struct InvalidDimension {
267 pub text: String,
269 pub expected: &'static str,
271}