Skip to main content

ferrin_spec/
image_model.rs

1//! Image model interface.
2
3use 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
23/// A model that generates images from a prompt.
24pub trait ImageModel: Send + Sync + 'static {
25    /// Provider identifier.
26    fn provider(&self) -> &ProviderId;
27
28    /// Model identifier.
29    fn model_id(&self) -> &ModelId;
30
31    /// Maximum number of images per call, or `None` when unknown (treated as 1).
32    fn max_images_per_call(&self) -> Option<usize>;
33
34    /// Generates `options.n` images.
35    fn do_generate(
36        &self,
37        options: ImageOptions,
38    ) -> impl Future<Output = Result<ImageResult, ProviderError>> + Send;
39}
40
41/// Options for an image generation call.
42#[derive(Debug, Clone)]
43pub struct ImageOptions {
44    /// Text prompt; `None` for pure edit/variation calls.
45    pub prompt: Option<String>,
46    /// Number of images to generate.
47    pub n: u32,
48    /// Requested size.
49    pub size: Option<ImageSize>,
50    /// Requested aspect ratio.
51    pub aspect_ratio: Option<AspectRatio>,
52    /// Random seed.
53    pub seed: Option<u64>,
54    /// Reference or input images.
55    pub files: Vec<ImageFile>,
56    /// Mask image for edits.
57    pub mask: Option<ImageFile>,
58    /// Provider-specific options keyed by provider name.
59    pub provider_options: ProviderOptions,
60    /// Additional request headers.
61    pub headers: Headers,
62    /// Cancellation token.
63    pub cancellation: CancellationToken,
64}
65
66impl ImageOptions {
67    /// Creates options for `prompt` requesting one image.
68    #[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/// An input image: inline bytes or a URL.
95#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
96pub struct ImageFile {
97    /// Image payload.
98    pub data: FileData,
99    /// Media type, if known.
100    #[serde(default, skip_serializing_if = "Option::is_none")]
101    pub media_type: Option<MediaType>,
102    /// Provider-specific options for this file.
103    #[serde(default, skip_serializing_if = "Option::is_none")]
104    pub provider_options: Option<ProviderOptions>,
105}
106
107/// A generated image.
108#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
109pub struct GeneratedImage {
110    /// Image bytes.
111    #[serde(with = "base64_bytes")]
112    pub data: Bytes,
113    /// Media type, if the provider reports it; otherwise detected by the core.
114    #[serde(default, skip_serializing_if = "Option::is_none")]
115    pub media_type: Option<MediaType>,
116}
117
118/// Token usage of an image call.
119#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
120pub struct ImageUsage {
121    /// Input tokens.
122    #[serde(default, skip_serializing_if = "Option::is_none")]
123    pub input_tokens: Option<u64>,
124    /// Output tokens.
125    #[serde(default, skip_serializing_if = "Option::is_none")]
126    pub output_tokens: Option<u64>,
127    /// Total tokens.
128    #[serde(default, skip_serializing_if = "Option::is_none")]
129    pub total_tokens: Option<u64>,
130}
131
132/// Result of an image generation call.
133#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
134pub struct ImageResult {
135    /// Generated images.
136    pub images: Vec<GeneratedImage>,
137    /// Whether an empty result may be retried; `None` means "use the default".
138    #[serde(default, skip_serializing_if = "Option::is_none")]
139    pub is_retryable: Option<bool>,
140    /// Warnings.
141    #[serde(default)]
142    pub warnings: Vec<Warning>,
143    /// Provider-specific metadata (per-image metadata under `images`).
144    #[serde(default, skip_serializing_if = "Option::is_none")]
145    pub provider_metadata: Option<ProviderMetadata>,
146    /// Response metadata; `timestamp` and `model_id` are expected to be set.
147    #[serde(default)]
148    pub response: ResponseMetadata,
149    /// Token usage, if reported.
150    #[serde(default, skip_serializing_if = "Option::is_none")]
151    pub usage: Option<ImageUsage>,
152}
153
154/// Image size in pixels, serialized as `WIDTHxHEIGHT`.
155#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
156#[serde(try_from = "String", into = "String")]
157pub struct ImageSize {
158    /// Width in pixels.
159    pub width: u32,
160    /// Height in pixels.
161    pub height: u32,
162}
163
164impl ImageSize {
165    /// Creates a size.
166    #[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/// Aspect ratio, serialized as `WIDTH:HEIGHT`.
206#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
207#[serde(try_from = "String", into = "String")]
208pub struct AspectRatio {
209    /// Horizontal component.
210    pub width: u32,
211    /// Vertical component.
212    pub height: u32,
213}
214
215impl AspectRatio {
216    /// Creates a ratio.
217    #[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/// Error returned when a size or aspect ratio string is malformed.
264#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
265#[error("invalid dimension `{text}`: expected `{expected}`")]
266pub struct InvalidDimension {
267    /// The rejected text.
268    pub text: String,
269    /// The expected format.
270    pub expected: &'static str,
271}