aither_openai/
image.rs

1use crate::{
2    client::{Config, OpenAI},
3    error::OpenAIError,
4};
5use aither_core::image::{Data, ImageGenerator, Prompt, Size};
6use base64::{Engine as _, engine::general_purpose};
7use futures_core::Stream;
8use futures_lite::StreamExt;
9use serde::{Deserialize, Serialize};
10use std::sync::Arc;
11use zenwave::{
12    Client, client, header,
13    multipart::{MultipartPart, encode as encode_multipart},
14};
15
16impl ImageGenerator for OpenAI {
17    type Error = OpenAIError;
18
19    fn create(
20        &self,
21        prompt: Prompt,
22        size: Size,
23    ) -> impl Stream<Item = Result<Data, Self::Error>> + Send {
24        let cfg = self.config();
25        let prompt_text = prompt.text().to_owned();
26        let size_token = format_size(size);
27        futures_lite::stream::iter(vec![generate_images(cfg, prompt_text, size_token)])
28            .then(|fut| fut)
29            .map(|result| result.map(futures_lite::stream::iter).ok())
30            .filter_map(core::convert::identity)
31            .flatten()
32            .map(Ok)
33    }
34
35    fn edit(
36        &self,
37        prompt: Prompt,
38        mask: &[u8],
39    ) -> impl Stream<Item = Result<Data, Self::Error>> + Send {
40        let cfg = self.config();
41        let prompt_text = prompt.text().to_owned();
42        let size_token = format_size(Size::square(1024));
43        let base = prompt.images().first().cloned().ok_or_else(|| {
44            OpenAIError::Api("image editing requires a base image via Prompt::with_image".into())
45        });
46        let mask_bytes = mask.to_vec();
47        futures_lite::stream::iter(base).flat_map(move |base_image| {
48            futures_lite::stream::iter(vec![edit_image(
49                cfg.clone(),
50                prompt_text.clone(),
51                size_token.clone(),
52                base_image,
53                mask_bytes.clone(),
54            )])
55            .then(|fut| fut)
56            .map(|result| result.map(futures_lite::stream::iter).ok())
57            .filter_map(core::convert::identity)
58            .flatten()
59            .map(Ok)
60        })
61    }
62}
63
64async fn generate_images(
65    cfg: Arc<Config>,
66    prompt: String,
67    size: String,
68) -> Result<Vec<Data>, OpenAIError> {
69    let endpoint = cfg.request_url("/images/generations");
70    let mut backend = client();
71    let mut builder = backend
72        .post(endpoint)
73        .map_err(OpenAIError::Http)?
74        .header(header::AUTHORIZATION.as_str(), cfg.request_auth())
75        .map_err(OpenAIError::Http)?
76        .header(header::USER_AGENT.as_str(), "aither-openai/0.1")
77        .map_err(OpenAIError::Http)?;
78    if let Some(org) = &cfg.organization {
79        builder = builder
80            .header("OpenAI-Organization", org.clone())
81            .map_err(OpenAIError::Http)?;
82    }
83    let request = ImageGenerationRequest {
84        model: &cfg.image_model,
85        prompt: &prompt,
86        size: &size,
87        response_format: "b64_json",
88        n: 1,
89    };
90    let response: ImageResponse = builder
91        .json_body(&request)
92        .map_err(OpenAIError::Http)?
93        .json()
94        .await
95        .map_err(OpenAIError::Http)?;
96    response.into_images()
97}
98
99async fn edit_image(
100    cfg: Arc<Config>,
101    prompt: String,
102    size: String,
103    image: Vec<u8>,
104    mask: Vec<u8>,
105) -> Result<Vec<Data>, OpenAIError> {
106    let endpoint = cfg.request_url("/images/edits");
107    let mut backend = client();
108    let mut builder = backend
109        .post(endpoint)
110        .map_err(OpenAIError::Http)?
111        .header(header::AUTHORIZATION.as_str(), cfg.request_auth())
112        .map_err(OpenAIError::Http)?
113        .header(header::USER_AGENT.as_str(), "aither-openai/0.1")
114        .map_err(OpenAIError::Http)?;
115    if let Some(org) = &cfg.organization {
116        builder = builder
117            .header("OpenAI-Organization", org.clone())
118            .map_err(OpenAIError::Http)?;
119    }
120    let mut parts = vec![
121        MultipartPart::text("model", cfg.image_model.clone()),
122        MultipartPart::text("prompt", prompt),
123        MultipartPart::text("size", size),
124        MultipartPart::text("response_format", "b64_json"),
125        MultipartPart::binary("image", "image.png", "application/octet-stream", image),
126    ];
127    if !mask.is_empty() {
128        parts.push(MultipartPart::binary(
129            "mask",
130            "mask.png",
131            "application/octet-stream",
132            mask,
133        ));
134    }
135    let (boundary, body) = encode_multipart(parts);
136    let response: ImageResponse = builder
137        .header(
138            header::CONTENT_TYPE.as_str(),
139            format!("multipart/form-data; boundary={boundary}"),
140        )
141        .map_err(OpenAIError::Http)?
142        .bytes_body(body)
143        .json()
144        .await
145        .map_err(OpenAIError::Http)?;
146    response.into_images()
147}
148
149fn format_size(size: Size) -> String {
150    format!("{}x{}", size.width(), size.height())
151}
152
153#[derive(Debug, Serialize)]
154struct ImageGenerationRequest<'a> {
155    model: &'a str,
156    prompt: &'a str,
157    size: &'a str,
158    #[serde(rename = "response_format")]
159    response_format: &'static str,
160    n: u8,
161}
162
163#[derive(Debug, Deserialize)]
164struct ImageResponse {
165    data: Vec<ImagePayload>,
166}
167
168impl ImageResponse {
169    fn into_images(self) -> Result<Vec<Data>, OpenAIError> {
170        self.data
171            .into_iter()
172            .map(ImagePayload::into_bytes)
173            .collect()
174    }
175}
176
177#[derive(Debug, Deserialize)]
178struct ImagePayload {
179    #[serde(default)]
180    b64_json: Option<String>,
181}
182
183impl ImagePayload {
184    fn into_bytes(self) -> Result<Data, OpenAIError> {
185        let encoded = self
186            .b64_json
187            .ok_or_else(|| OpenAIError::Api("image response missing `b64_json` field".into()))?;
188        general_purpose::STANDARD
189            .decode(encoded)
190            .map_err(OpenAIError::from)
191    }
192}
aither_openai/image.rs

aither_openai/
image.rs
