rig_bedrock/
image.rs

1use crate::client::Client;
2use crate::types::errors::AwsSdkInvokeModelError;
3use crate::types::text_to_image::{TextToImageGeneration, TextToImageResponse};
4use aws_smithy_types::Blob;
5use rig_core::image_generation::{
6    self, ImageGenerationError, ImageGenerationRequest, ImageGenerationResponse,
7};
8
9// The model-id string values are canonically defined in `crate::completion`.
10// The Titan image generators are gone: `amazon.titan-image-generator-v1` and
11// `-v2:0` are absent from `ListFoundationModels` in every region checked
12// (us-east-1, us-west-2, eu-central-1, ap-northeast-1), so their aliases are
13// removed rather than left pointing at identifiers Bedrock rejects.
14pub use crate::completion::{
15    AMAZON_NOVA_CANVAS, STABILITY_SD3_5_LARGE, STABILITY_STABLE_IMAGE_CORE_1_0,
16    STABILITY_STABLE_IMAGE_ULTRA_1_0,
17};
18
19#[derive(Clone)]
20pub struct ImageGenerationModel {
21    pub(crate) client: Client,
22    pub model: String,
23}
24
25impl ImageGenerationModel {
26    pub fn new(client: Client, model: impl Into<String>) -> Self {
27        Self {
28            client,
29            model: model.into(),
30        }
31    }
32}
33
34impl image_generation::ImageGenerationModel for ImageGenerationModel {
35    type Response = TextToImageResponse;
36
37    type Client = Client;
38
39    fn make(client: &Self::Client, model: impl Into<String>) -> Self {
40        Self::new(client.clone(), model)
41    }
42
43    async fn image_generation(
44        &self,
45        generation_request: ImageGenerationRequest,
46    ) -> Result<ImageGenerationResponse<Self::Response>, ImageGenerationError> {
47        let mut request = TextToImageGeneration::new(generation_request.prompt);
48        request.width(generation_request.width);
49        request.height(generation_request.height);
50
51        let body = serde_json::to_string(&request)?;
52        let model_response = self
53            .client
54            .get_inner()
55            .await
56            .invoke_model()
57            .model_id(self.model.as_str())
58            .content_type("application/json")
59            .accept("application/json")
60            .body(Blob::new(body))
61            .send()
62            .await
63            .map_err(|sdk_error| {
64                Into::<ImageGenerationError>::into(AwsSdkInvokeModelError(sdk_error))
65            })?;
66
67        let response_str = String::from_utf8(model_response.body.into_inner())
68            .map_err(|e| ImageGenerationError::ResponseError(e.to_string()))?;
69
70        let result: TextToImageResponse = serde_json::from_str(&response_str)
71            .map_err(|e| ImageGenerationError::ResponseError(e.to_string()))?;
72
73        result.try_into()
74    }
75}
rig_bedrock/image.rs

rig_bedrock/
image.rs
