Как генерировать изображения на Rust в 100 строчек кода с Yandex ART?

Если в вашем Rust-проекте возникает необходимость генерировать изображения, то расскажите зачем) А о том, как это сделать — в этой статье. В качестве источника самих картинок я выбрал Yandex ART из-за того, что с ним не нужно возиться со всякими трехбуквенными сервисами, реклама которых в России запрещена.

Итак, для генерации будем использовать Yandex ART, документация по нему представлена здесь. Но напрямую работать с API мы не будем, а воспользуемся готовой библиотекой nn_yandex_art версии 0.2.0: Crates.io.

Github-репозиторий nn_yandex_art: Click

Какие нужны предварительные приготовления?

1. Создать проект

2. Скачать библиотеку

3. Получить API-ключ к Yandex ART

Что делаем дальше? Импортируем нужные структуры:

use nn_yandex_art::Art;
use nn_yandex_art::models::request::message::MessageBuilder;
use nn_yandex_art::models::request::aspect_ratio::AspectRatioBuilder;
use nn_yandex_art::models::request::generation_options::GenerationOptionsBuilder;
use nn_yandex_art::models::request::types::ImageType;
use nn_yandex_art::models::request::RequestBuilder;

Кроме этого, нам понадобится ряд других библиотек:

use anyhow;
use tokio::time::{sleep, Duration};
use std::fs::File;
use std::io::Write;
use base64::engine::general_purpose::STANDARD;
use base64::Engine;
use std::env;

В итоге Cargo.toml будет выглядеть следующим образом:

[dependencies]
nn_yandex_art = "0.2.0"
anyhow = "1.0.100"
tokio = { version = "1.47.1", features = ["rt", "rt-multi-thread", "macros"] }
base64 = "0.22.1"
dotenvy = "0.15.7"

После этого уже приступаем непосредственно к написанию генератора изображений. У него есть несколько важных настроек: текстовый промт, соотношение сторон, а также тип изображения (jpg или png). Сам генератор помещается буквально в две функции.

pub async fn generate_image(prompt: &str, path: &str, width_ratio: i64, height_ratio: i64) -> Result<(), anyhow::Error>{

    let BUCKET = env::var("BUCKET")?;
    let API_KEY = env::var("API")?;

    let message = MessageBuilder::new()
        .text(prompt)
        .weight(1)
        .build()?;


    let aspect_ratio = AspectRatioBuilder::new()
        .width_ratio(width_ratio)
        .height_ratio(height_ratio)
        .build();

    let generation_options = GenerationOptionsBuilder::new()
        .aspect_ratio(aspect_ratio)
        .mime_type(ImageType::Png)
        //.seed(121212121212) // !Optional
        .build()?;

    let request = RequestBuilder::new()
        .generation_options(generation_options)
        .message(message)
        .build()?;

    let art = Art::new(API_KEY, BUCKET);
    let mut res = art.generate_image(request).await?;
    let id = res.id;

    if let Some(e) = res.error{
        return Err(anyhow::anyhow!("{}", e.message))
    }

    while !res.done{
        sleep(Duration::from_secs(1)).await;
        res = art.check_operation(&id).await?
    }

    if let Some(resp) = res.response {
        save_image(resp.image, path)
    } else {
        Err(anyhow::anyhow!("Response is missing image data"))
    }
}

fn save_image(image: String, path: &str) -> Result<(), anyhow::Error>{
    match STANDARD.decode(image) {
        Ok(bytes) => {
            match File::create(path) {
                Ok(mut file) => {
                    if let Err(e) = file.write_all(&bytes) {
                        Err(anyhow::anyhow!("Error writing file: {e}"))
                    } else {
                        Ok(())
                    }
                }
                Err(e) => Err(anyhow::anyhow!(e)),
            }
        }
        Err(e) => Err(anyhow::anyhow!(e)),
    }
}

В данном случае сгенерированные изображения будут сохраняться в файлы. Текстовый промт, соотношение сторон и сам путь к файлу передаются в функцию generate_image(prompt: &str, path: &str, width_ratio: i64, height_ratio: i64), а функция save_image(image: String, path: &str) непосредственно сохраняет сгенерированный BASE64 в файл.

Далее функция main:

#[tokio::main]
async fn main() {
    dotenv().ok();
    let prompt = "Кошка сидит на диване";
    let path = "./images/image4.png";
    let width_ratio: i64 = 3;
    let height_ratio: i64 = 2;
    match generate_image(prompt, path, width_ratio, height_ratio).await{
        Ok(_) => println!("Image generated successfully"),
        Err(e) => eprintln!("{e}"),
    }
}

И на этом все. Генератор изображений готов. Полный код проекта можно найти здесь. Проведя эксперименты с различными промтами на русском и английском языках я получил следующие результаты: Кликни для просмотра.