1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

use reqwest::Body;
use tokio_util::codec::{BytesCodec, FramedRead};

use crate::{
    error::OpenAIError,
    types::{
        CreateImageEditRequest, CreateImageRequest, CreateImageVariationRequest, ImageInput,
        ImageResponse,
    },
    Client,
};

/// Given a prompt and/or an input image, the model will generate a new image.
///
/// Related guide: [Image generation](https://beta.openai.com/docs/guides/images/introduction)
pub struct Image;

impl Image {
    /// Creates an image given a prompt.
    pub async fn create(
        client: &Client,
        request: CreateImageRequest,
    ) -> Result<ImageResponse, OpenAIError> {
        client.post("/images/generations", request).await
    }

    pub(crate) async fn file_stream_body(image_input: &ImageInput) -> Result<Body, OpenAIError> {
        let file = tokio::fs::File::open(image_input.path.as_path())
            .await
            .map_err(|e| OpenAIError::ImageReadError(e.to_string()))?;
        let stream = FramedRead::new(file, BytesCodec::new());
        let body = Body::wrap_stream(stream);
        Ok(body)
    }

    /// Creates the part for the given image file for multipart upload.
    pub(crate) async fn create_part(
        image_input: &ImageInput,
    ) -> Result<reqwest::multipart::Part, OpenAIError> {
        let image_file_name = image_input
            .path
            .as_path()
            .file_name()
            .ok_or_else(|| {
                OpenAIError::ImageReadError(format!(
                    "cannot extract file name from {:#?}",
                    image_input.path
                ))
            })?
            .to_str()
            .unwrap()
            .to_string();

        let image_part =
            reqwest::multipart::Part::stream(Image::file_stream_body(image_input).await?)
                .file_name(image_file_name)
                .mime_str("application/octet-stream")
                .unwrap();

        Ok(image_part)
    }

    /// Creates an edited or extended image given an original image and a prompt.
    pub async fn create_edit(
        client: &Client,
        request: CreateImageEditRequest,
    ) -> Result<ImageResponse, OpenAIError> {
        let image_part = Image::create_part(&request.image).await?;
        let mask_part = Image::create_part(&request.mask).await?;

        let mut form = reqwest::multipart::Form::new()
            .part("image", image_part)
            .part("mask", mask_part)
            .text("prompt", request.prompt);

        if request.n.is_some() {
            form = form.text("n", request.n.unwrap().to_string())
        }

        if request.size.is_some() {
            form = form.text("size", request.size.unwrap().to_string())
        }

        if request.response_format.is_some() {
            form = form.text(
                "response_format",
                request.response_format.unwrap().to_string(),
            )
        }

        if request.user.is_some() {
            form = form.text("user", request.user.unwrap())
        }

        client.post_form("/images/edits", form).await
    }

    /// Creates a variation of a given image.
    pub async fn create_variation(
        client: &Client,
        request: CreateImageVariationRequest,
    ) -> Result<ImageResponse, OpenAIError> {
        let image_part = Image::create_part(&request.image).await?;

        let mut form = reqwest::multipart::Form::new().part("image", image_part);

        if request.n.is_some() {
            form = form.text("n", request.n.unwrap().to_string())
        }

        if request.size.is_some() {
            form = form.text("size", request.size.unwrap().to_string())
        }

        if request.response_format.is_some() {
            form = form.text(
                "response_format",
                request.response_format.unwrap().to_string(),
            )
        }

        if request.user.is_some() {
            form = form.text("user", request.user.unwrap())
        }

        client.post_form("/images/variations", form).await
    }
}