rig_core/providers/huggingface/
image_generation.rs

1use super::client::Client;
2use crate::http_client::HttpClientExt;
3use crate::image_generation;
4use crate::image_generation::{ImageGenerationError, ImageGenerationRequest};
5use serde_json::json;
6
7#[allow(non_upper_case_globals)]
8pub mod image_generation_models {
9    pub const Flux1: &str = "black-forest-labs/FLUX.1-dev";
10    pub const Kolors: &str = "Kwai-Kolors/Kolors";
11    pub const StableDiffusion3: &str = "stabilityai/stable-diffusion-3-medium-diffusers";
12}
13pub use image_generation_models::*;
14
15#[derive(Debug)]
16pub struct ImageGenerationResponse {
17    data: Vec<u8>,
18}
19
20impl TryFrom<ImageGenerationResponse>
21    for image_generation::ImageGenerationResponse<ImageGenerationResponse>
22{
23    type Error = ImageGenerationError;
24
25    fn try_from(value: ImageGenerationResponse) -> Result<Self, Self::Error> {
26        Ok(image_generation::ImageGenerationResponse {
27            image: value.data.clone(),
28            response: value,
29        })
30    }
31}
32
33#[derive(Clone)]
34pub struct ImageGenerationModel<T = reqwest::Client> {
35    client: Client<T>,
36    pub model: String,
37}
38
39impl<T> ImageGenerationModel<T> {
40    pub fn new(client: Client<T>, model: impl Into<String>) -> Self {
41        ImageGenerationModel {
42            client,
43            model: model.into(),
44        }
45    }
46}
47
48impl<T> image_generation::ImageGenerationModel for ImageGenerationModel<T>
49where
50    T: HttpClientExt + Send + Clone + 'static,
51{
52    type Response = ImageGenerationResponse;
53
54    type Client = Client<T>;
55
56    fn make(client: &Self::Client, model: impl Into<String>) -> Self {
57        Self::new(client.clone(), model)
58    }
59
60    async fn image_generation(
61        &self,
62        request: ImageGenerationRequest,
63    ) -> Result<image_generation::ImageGenerationResponse<Self::Response>, ImageGenerationError>
64    {
65        let request = json!({
66            "inputs": request.prompt,
67            "parameters": {
68                "width": request.width,
69                "height": request.height
70            }
71        });
72
73        let route = self
74            .client
75            .subprovider()
76            .image_generation_endpoint(&self.model)?;
77
78        let body = serde_json::to_vec(&request)?;
79
80        let req = self
81            .client
82            .post(&route)?
83            .header("Content-Type", "application/json")
84            .body(body)
85            .map_err(|e| ImageGenerationError::HttpError(e.into()))?;
86
87        let response = self.client.send(req).await?;
88
89        if !response.status().is_success() {
90            let status = response.status();
91            let text: Vec<u8> = response.into_body().await?;
92
93            return Err(ImageGenerationError::from_http_response(
94                status,
95                String::from_utf8_lossy(&text),
96            ));
97        }
98
99        let data: Vec<u8> = response.into_body().await?;
100
101        ImageGenerationResponse { data }.try_into()
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108    use crate::client::image_generation::ImageGenerationClient;
109    use crate::image_generation::ImageGenerationModel as _;
110    use crate::test_utils::RecordingHttpClient;
111
112    #[tokio::test]
113    async fn image_generation_non_success_preserves_status_and_body() {
114        let body = r#"{"error":{"message":"boom"}}"#;
115        let http_client =
116            RecordingHttpClient::with_error_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
117        let client = Client::builder()
118            .api_key("test-key")
119            .http_client(http_client)
120            .build()
121            .expect("build client");
122        let model = client.image_generation_model(Flux1);
123
124        let request = model
125            .image_generation_request()
126            .prompt("draw a cat")
127            .build();
128
129        let error = model
130            .image_generation(request)
131            .await
132            .expect_err("should fail with non-success status");
133
134        assert!(matches!(error, ImageGenerationError::HttpError(_)));
135        assert_eq!(
136            error.provider_response_status(),
137            Some(http::StatusCode::SERVICE_UNAVAILABLE)
138        );
139        assert_eq!(error.provider_response_body(), Some(body));
140    }
141}
rig_core/providers/huggingface/image_generation.rs

rig_core/providers/huggingface/
image_generation.rs
