aither_core/
image.rs

1use alloc::{boxed::Box, string::String, sync::Arc, vec::Vec};
2use futures_core::Stream;
3
4/// Image data as bytes.
5///
6/// Type alias for [`Vec<u8>`] representing image data.
7pub type Data = Vec<u8>;
8
9/// Trait for generating and editing images from prompts and masks.
10///
11/// Images are returned as a stream where each item represents a complete image
12/// with progressively improving quality, allowing for real-time preview during generation.
13pub trait ImageGenerator {
14    /// The error type returned by the image generator.
15    type Error: core::error::Error + Send + Sync + 'static;
16
17    /// Create an image from a prompt and a specified size.
18    ///
19    /// # Arguments
20    ///
21    /// * `prompt` - The prompt containing text and optional images.
22    /// * `size` - The desired size of the generated image.
23    ///
24    /// # Returns
25    ///
26    /// A stream where each item is a complete image with progressively improving quality.
27    fn create(
28        &self,
29        prompt: Prompt,
30        size: Size,
31    ) -> impl Stream<Item = Result<Data, Self::Error>> + Send;
32
33    /// Edit an image using a prompt and a mask.
34    ///
35    /// # Arguments
36    ///
37    /// * `prompt` - The prompt containing text and optional images.
38    /// * `mask` - The mask to apply to the image data.
39    ///
40    /// # Returns
41    ///
42    /// A stream where each item is a complete image with progressively improving quality.
43    fn edit(
44        &self,
45        prompt: Prompt,
46        mask: &[u8],
47    ) -> impl Stream<Item = Result<Data, Self::Error>> + Send;
48}
49
50macro_rules! impl_image_generator {
51    ($($name:ident),*) => {
52        $(
53            impl<T: ImageGenerator> ImageGenerator for $name<T> {
54                type Error = T::Error;
55
56                fn create(
57                    &self,
58                    prompt: Prompt,
59                    size: Size,
60                ) -> impl Stream<Item = Result<Data, Self::Error>> + Send {
61                    T::create(self, prompt, size)
62                }
63
64                fn edit(
65                    &self,
66                    prompt: Prompt,
67                    mask: &[u8],
68                ) -> impl Stream<Item = Result<Data, Self::Error>> + Send {
69                    T::edit(self, prompt, mask)
70                }
71            }
72        )*
73    };
74}
75
76impl_image_generator!(Arc, Box);
77
78/// Represents a prompt for image generation, including text and optional images.
79#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
80#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
81pub struct Prompt {
82    /// The text description for the image generation.
83    text: String,
84    /// Optional images to guide the generation process.
85    image: Vec<Data>,
86}
87
88impl Prompt {
89    /// Creates a new `Prompt` with the given text
90    #[must_use]
91    pub fn new(text: impl Into<String>) -> Self {
92        Self {
93            text: text.into(),
94            image: Vec::new(),
95        }
96    }
97
98    /// Returns the text description of the prompt.
99    #[must_use]
100    pub fn text(&self) -> &str {
101        &self.text
102    }
103
104    /// Returns the images associated with the prompt.
105    #[must_use]
106    pub fn images(&self) -> &[Data] {
107        &self.image
108    }
109
110    /// Adds an image to the prompt and returns the updated `Prompt`.
111    ///
112    /// # Arguments
113    ///
114    /// * `image` - The image data to add to the prompt.
115    #[must_use]
116    pub fn with_image(mut self, image: Data) -> Self {
117        self.image.push(image);
118        self
119    }
120}
121
122impl From<String> for Prompt {
123    /// Converts a `String` into a `Prompt`.
124    ///
125    /// # Arguments
126    ///
127    /// * `text` - The text to use for the prompt.
128    fn from(text: String) -> Self {
129        Self::new(text)
130    }
131}
132
133impl From<&str> for Prompt {
134    /// Converts a `&str` into a `Prompt`.
135    ///
136    /// # Arguments
137    ///
138    /// * `text` - The text to use for the prompt.
139    fn from(text: &str) -> Self {
140        Self::new(text)
141    }
142}
143
144/// Represents the size (width and height) of an image.
145#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
146#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
147pub struct Size {
148    /// The width of the image in pixels.
149    width: u32,
150    /// The height of the image in pixels.
151    height: u32,
152}
153
154impl Size {
155    /// Creates a new `Size` with the given width and height.
156    ///
157    /// # Arguments
158    ///
159    /// * `width` - The width of the image in pixels.
160    /// * `height` - The height of the image in pixels.
161    #[must_use]
162    pub const fn new(width: u32, height: u32) -> Self {
163        Self { width, height }
164    }
165
166    /// Creates a new `Size` with equal width and height (a square).
167    ///
168    /// # Arguments
169    ///
170    /// * `size` - The width and height of the square image in pixels.
171    #[must_use]
172    pub const fn square(size: u32) -> Self {
173        Self {
174            width: size,
175            height: size,
176        }
177    }
178
179    /// Returns the width of the image in pixels.
180    #[must_use]
181    pub const fn width(&self) -> u32 {
182        self.width
183    }
184
185    /// Returns the height of the image in pixels.
186    #[must_use]
187    pub const fn height(&self) -> u32 {
188        self.height
189    }
190
191    /// Returns the total number of pixels (width × height).
192    #[must_use]
193    pub const fn pixel_count(&self) -> u64 {
194        self.width as u64 * self.height as u64
195    }
196
197    /// Returns whether this is a square image (width equals height).
198    #[must_use]
199    pub const fn is_square(&self) -> bool {
200        self.width == self.height
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use core::convert::Infallible;
207
208    use super::*;
209    use alloc::vec;
210    use futures_lite::StreamExt;
211
212    struct MockImageGenerator;
213
214    impl ImageGenerator for MockImageGenerator {
215        type Error = Infallible;
216        fn create(
217            &self,
218            prompt: Prompt,
219            _size: Size,
220        ) -> impl Stream<Item = Result<Data, Self::Error>> + Send {
221            // Create mock image data based on prompt
222            let prompt_bytes = prompt.text.as_bytes();
223            let chunk1 = prompt_bytes.to_vec();
224            let chunk2 = vec![0xFF, 0xD8, 0xFF, 0xE0]; // Mock JPEG header
225            let chunk3 = vec![0x00; 100]; // Mock image data
226
227            futures_lite::stream::iter(vec![chunk1, chunk2, chunk3].into_iter().map(Ok))
228        }
229
230        fn edit(
231            &self,
232            prompt: Prompt,
233            _mask: &[u8],
234        ) -> impl Stream<Item = Result<Data, Self::Error>> + Send {
235            // Create mock image data based on prompt
236            let prompt_bytes = prompt.text.as_bytes();
237            let chunk1 = prompt_bytes.to_vec();
238            let chunk2 = vec![0xFF, 0xD8, 0xFF, 0xE0]; // Mock JPEG header
239            let chunk3 = vec![0x00; 100]; // Mock image data
240
241            futures_lite::stream::iter(vec![chunk1, chunk2, chunk3].into_iter().map(Ok))
242        }
243    }
244
245    #[tokio::test]
246    async fn image_generation() {
247        let generator = MockImageGenerator;
248        let mut stream = generator.create(Prompt::new("a cat"), Size::square(256));
249
250        let mut chunks = Vec::new();
251        while let Some(chunk) = stream.next().await {
252            chunks.push(chunk.unwrap());
253        }
254
255        assert_eq!(chunks.len(), 3);
256        assert_eq!(chunks[0], b"a cat".to_vec());
257        assert_eq!(chunks[1], vec![0xFF, 0xD8, 0xFF, 0xE0]);
258        assert_eq!(chunks[2], vec![0x00; 100]);
259    }
260
261    #[tokio::test]
262    async fn image_generation_empty_prompt() {
263        let generator = MockImageGenerator;
264        let mut stream = generator.create(Prompt::new(""), Size::square(256));
265
266        let mut chunks = Vec::new();
267        while let Some(chunk) = stream.next().await {
268            chunks.push(chunk.unwrap());
269        }
270
271        assert_eq!(chunks.len(), 3);
272        assert_eq!(chunks[0], b"".to_vec());
273        assert_eq!(chunks[1], vec![0xFF, 0xD8, 0xFF, 0xE0]);
274        assert_eq!(chunks[2], vec![0x00; 100]);
275    }
276
277    #[tokio::test]
278    async fn image_generation_long_prompt() {
279        let generator = MockImageGenerator;
280        let long_prompt = "a very detailed and elaborate description of a beautiful landscape with mountains, rivers, and forests";
281        let mut stream = generator.create(Prompt::new(long_prompt), Size::square(512));
282
283        let mut total_bytes = 0;
284        while let Some(chunk) = stream.next().await {
285            total_bytes += chunk.unwrap().len();
286        }
287
288        // Should have prompt bytes + header bytes + 100 mock data bytes
289        assert_eq!(total_bytes, long_prompt.len() + 4 + 100);
290    }
291
292    #[tokio::test]
293    async fn data_type_alias() {
294        let data: Data = vec![1, 2, 3, 4];
295        assert_eq!(data.len(), 4);
296        assert_eq!(data[0], 1);
297        assert_eq!(data[3], 4);
298    }
299
300    #[test]
301    fn data_operations() {
302        let mut data: Data = vec![0xFF; 1024];
303        assert_eq!(data.len(), 1024);
304
305        data.push(0x00);
306        assert_eq!(data.len(), 1025);
307        assert_eq!(data[1024], 0x00);
308
309        data.extend_from_slice(&[0x01, 0x02]);
310        assert_eq!(data.len(), 1027);
311        assert_eq!(data[1025], 0x01);
312        assert_eq!(data[1026], 0x02);
313    }
314}
aither_core/image.rs

aither_core/
image.rs
