Add text to image

Author: rikhuijzer

.gitignore

@@ -1,6 +1,7 @@
 **/target/
 **/*.rs.bk
-ata.toml
 Cargo.lock
 test.env
-**/tmp*
+**/tmp*
+**/*.png
+**/*.jpg

README.md

@@ -12,6 +12,19 @@ $ export OPENAI_KEY="$(grep 'OPENAI_KEY' .env | cut -d= -f2)"
 $ cat myfile.txt | trf tts | vlc - --intf dummy
 ```
 
+Or generate an image from text:
+
+```sh
+$ export DEEPINFRA_KEY="$(grep 'DEEPINFRA_KEY' .env | cut -d= -f2)"
+
+$ echo "A photo of a beach in Hawaii" | cargo run -- --verbose tti \
+    --model=black-forest-labs/FLUX-1-dev \
+    --output=myfile --width=1024 --height=1024 --steps=25
+
+$ ls myfile*
+myfile.png
+```
+
 You can also chat with an LLM:
 
 ```sh

src/main.rs

@@ -1,11 +1,13 @@
 mod chat;
+mod tti;
 mod tts;
 
 use chat::ChatArgs;
 use clap::Parser;
 use std::io::Read;
 use tracing::subscriber::SetGlobalDefaultError;
 use transformrs::Key;
+use tti::TextToImageArgs;
 use tts::TextToSpeechArgs;
 
 #[derive(clap::Subcommand)]
@@ -15,6 +17,11 @@ enum Commands {
     /// Takes text input from stdin and chats with an AI model.
     #[command()]
     Chat(ChatArgs),
+    /// Convert text to image
+    ///
+    /// Takes text input from stdin and converts it to an image using text-to-image models.
+    #[command()]
+    Tti(TextToImageArgs),
     /// Convert text to speech
     ///
     /// Takes text input from stdin and converts it to speech using text-to-speech models.
@@ -83,6 +90,9 @@ async fn main() {
         Commands::Chat(args) => {
             chat::chat(&args, &key, &input).await;
         }
+        Commands::Tti(args) => {
+            tti::tti(&args, &key, &input).await;
+        }
         Commands::Tts(args) => {
             tts::tts(&args, &key, &input).await;
         }

src/tti.rs

@@ -0,0 +1,66 @@
+use std::fs::File;
+use std::io::Write;
+use transformrs::text_to_image::TTIConfig;
+use transformrs::Provider;
+
+#[derive(clap::Parser)]
+pub(crate) struct TextToImageArgs {
+    /// Model to use (optional)
+    #[arg(long)]
+    model: Option<String>,
+
+    /// Number of steps (optional)
+    #[arg(long, default_value_t = 10)]
+    steps: u32,
+
+    /// CFG scale (optional)
+    #[arg(long, default_value_t = 3)]
+    cfg_scale: u32,
+
+    /// Height (optional)
+    #[arg(long, default_value_t = 512)]
+    height: u32,
+
+    /// Width (optional)
+    #[arg(long, default_value_t = 512)]
+    width: u32,
+
+    /// Output filename without extension
+    #[arg(long, short = 'o')]
+    output: Option<String>,
+}
+
+fn default_model(provider: &Provider) -> String {
+    match provider {
+        Provider::DeepInfra => "black-forest-labs/FLUX-1-schnell".to_string(),
+        _ => "black-forest-labs/FLUX-1-dev".to_string(),
+    }
+}
+
+pub(crate) async fn tti(args: &TextToImageArgs, key: &transformrs::Key, input: &str) {
+    let provider = key.provider.clone();
+    let config = TTIConfig {
+        model: args
+            .model
+            .clone()
+            .unwrap_or_else(|| default_model(&provider)),
+        steps: Some(args.steps),
+        cfg_scale: Some(args.cfg_scale),
+        height: Some(args.height),
+        width: Some(args.width),
+    };
+    let resp = transformrs::text_to_image::text_to_image(key, config, input)
+        .await
+        .unwrap()
+        .structured()
+        .unwrap();
+    let encoded = &resp.images[0];
+    let image = encoded.base64_decode().unwrap();
+    if let Some(output) = &args.output {
+        let filename = format!("{}.{}", output, image.filetype);
+        let mut file = File::create(filename).unwrap();
+        file.write_all(&image.image).unwrap();
+    } else {
+        std::io::stdout().write_all(&image.image).unwrap();
+    }
+}

tests/tti.rs

@@ -0,0 +1,51 @@
+mod common;
+
+use common::load_key;
+use common::trf;
+use transformrs::Provider;
+
+#[test]
+fn no_args() -> Result<(), Box<dyn std::error::Error>> {
+    let dir = tempfile::tempdir().unwrap();
+    let mut cmd = trf();
+    let key = load_key(&Provider::DeepInfra);
+    let cmd = cmd
+        .arg("tti")
+        .env("DEEPINFRA_KEY", key)
+        .write_stdin("Hello world")
+        .current_dir(&dir);
+    let output = cmd.assert().success().get_output().stdout.clone();
+
+    assert!(output.len() > 0);
+
+    Ok(())
+}
+
+fn default_settings_helper(provider: &Provider) -> Result<(), Box<dyn std::error::Error>> {
+    let dir = tempfile::tempdir().unwrap();
+    let mut cmd = trf();
+    let key = load_key(provider);
+    let name = provider.key_name();
+    cmd.arg("--verbose")
+        .arg("tti")
+        .arg("--output=myfile")
+        .arg("--width=128")
+        .arg("--height=128")
+        .arg("--steps=10")
+        .arg("--cfg-scale=3")
+        .env(name, key)
+        .write_stdin("image of a beach")
+        .current_dir(&dir)
+        .assert()
+        .success();
+
+    let path = dir.path().join("myfile.png");
+    assert!(path.exists());
+
+    Ok(())
+}
+
+#[test]
+fn default_settings_deepinfra() -> Result<(), Box<dyn std::error::Error>> {
+    default_settings_helper(&Provider::DeepInfra)
+}