comfyui_api/api/
upload.rs

1use reqwest::{multipart, Url};
2use serde::{Deserialize, Serialize};
3
4/// Errors that can occur when interacting with `UploadApi`.
5#[derive(thiserror::Error, Debug)]
6#[non_exhaustive]
7pub enum UploadApiError {
8    /// Error parsing endpoint URL
9    #[error("Failed to parse endpoint URL")]
10    ParseError(#[from] url::ParseError),
11    /// Error setting MIME type
12    #[error("Failed to set MIME type")]
13    SetMimeStrFailed(#[source] reqwest::Error),
14    /// Error sending request
15    #[error("Failed to send request")]
16    RequestFailed(#[from] reqwest::Error),
17    /// An error occurred while parsing the response from the API.
18    #[error("Parsing response failed")]
19    InvalidResponse(#[source] reqwest::Error),
20    /// An error occurred getting response data.
21    #[error("Failed to get response data")]
22    GetDataFailed(#[source] reqwest::Error),
23    /// Server returned an error when uploading file
24    #[error("Failed to upload image: {status}: {error}")]
25    UploadImageFailed {
26        status: reqwest::StatusCode,
27        error: String,
28    },
29}
30
31type Result<T> = std::result::Result<T, UploadApiError>;
32
33/// Struct representing an image.
34#[derive(Serialize, Deserialize, Debug, Clone)]
35pub struct ImageUpload {
36    /// The filename of the image.
37    pub name: String,
38    /// The subfolder.
39    pub subfolder: String,
40    /// The folder type.
41    #[serde(rename = "type")]
42    pub folder_type: String,
43}
44
45/// Struct representing a connection to the ComfyUI API `upload` endpoint.
46#[derive(Clone, Debug)]
47pub struct UploadApi {
48    client: reqwest::Client,
49    endpoint: Url,
50}
51
52impl UploadApi {
53    /// Constructs a new `UploadApi` client with a given `reqwest::Client` and ComfyUI API
54    /// endpoint.
55    ///
56    /// # Arguments
57    ///
58    /// * `client` - A `reqwest::Client` used to send requests.
59    /// * `endpoint` - A `str` representation of the endpoint url.
60    ///
61    /// # Returns
62    ///
63    /// A `Result` containing a new `UploadApi` instance on success, or an error if url parsing failed.
64    pub fn new<S>(client: reqwest::Client, endpoint: S) -> Result<Self>
65    where
66        S: AsRef<str>,
67    {
68        Ok(Self::new_with_url(client, Url::parse(endpoint.as_ref())?))
69    }
70
71    /// Constructs a new `UploadApi` client with a given `reqwest::Client` and endpoint `Url`.
72    ///
73    /// # Arguments
74    ///
75    /// * `client` - A `reqwest::Client` used to send requests.
76    /// * `endpoint` - A `Url` representing the endpoint url.
77    ///
78    /// # Returns
79    ///
80    /// A new `UploadApi` instance.
81    pub fn new_with_url(client: reqwest::Client, endpoint: Url) -> Self {
82        Self { client, endpoint }
83    }
84
85    /// Uploads an image using the `UploadApi` client.
86    ///
87    /// # Arguments
88    ///
89    /// * `image` - A `Vec<u8>` containing the image to upload.
90    ///
91    /// # Returns
92    ///
93    /// A `Result` containing an `Image` struct containing information about the image.
94    /// success, or an error if the request failed.
95    pub async fn image(&self, image: Vec<u8>) -> Result<ImageUpload> {
96        let file = multipart::Part::bytes(image)
97            .file_name("image.png")
98            .mime_str("image/png")
99            .map_err(UploadApiError::SetMimeStrFailed)?;
100        let form = multipart::Form::new().part("image", file);
101        let response = self
102            .client
103            .post(self.endpoint.clone().join("image")?)
104            .multipart(form)
105            .send()
106            .await
107            .map_err(UploadApiError::RequestFailed)?;
108        if response.status().is_success() {
109            return response
110                .json()
111                .await
112                .map_err(UploadApiError::InvalidResponse);
113        }
114        let status = response.status();
115        let text = response
116            .text()
117            .await
118            .map_err(UploadApiError::GetDataFailed)?;
119        Err(UploadApiError::UploadImageFailed {
120            status,
121            error: text,
122        })
123    }
124}