wip, don't use filesystem for image intermediates

crates/desktop/src-tauri/src/endpoints/image_endpoint.rs

@@ -1,13 +1,13 @@
 use crate::endpoints::sd::{run, Args};
-use crate::state::yaml_config::YamlConfig;
 use crate::ml::model_cache::ModelCache;
 use crate::ml::prompt_cache::PromptCache;
 use crate::state::app_config::AppConfig;
+use crate::state::yaml_config::YamlConfig;
 use anyhow::Result;
 use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
 use image::imageops::FilterType;
-use image::ImageReader;
-use log::error;
+use image::{DynamicImage, EncodableLayout, ImageReader};
+use log::{error, info};
 use std::io::Cursor;
 use std::path::PathBuf;
 use tauri::State;
@@ -46,13 +46,7 @@ pub fn infer_image(
   // TODO(bt,2025-02-17): Running out of vram with full image buffer size
   let image = image.resize(512, 512, FilterType::CatmullRom);
   
-  // TODO(bt,2025-02-14): Use byte buffers instead of hitting the filesystem.
-  let image_path = PathBuf::from("input_image.png");
-  
-  image.save(&image_path)
-    .map_err(|err| format!("Failed to save image: {}", err))?;
-  
-  let result = do_infer_image(&prompt, &image_path, &model_config, &model_cache, prompt_cache);
+  let result = do_infer_image(&prompt, image, &model_config, &model_cache, prompt_cache);
   
   if let Err(err) = result.as_deref() {
     error!("There was an error: {:?}", err);
@@ -63,7 +57,7 @@ pub fn infer_image(
 
 fn do_infer_image(
   prompt: &str,
-  image_path: &PathBuf,
+  image: DynamicImage,
   config: &AppConfig,
   model_cache: &ModelCache,
   prompt_cache: State<PromptCache>,
@@ -71,7 +65,7 @@ fn do_infer_image(
   println!("infer_image called; generating image with SDXL Turbo...");
   
   let args = Args {
-    image_path,
+    image: &image,
     prompt: prompt.to_string(),
     uncond_prompt: "".to_string(),
     //guidance_scale: Some(0.0),
@@ -82,17 +76,27 @@ fn do_infer_image(
   };
 
   match run(args) {
-    Ok(_) => {
+    Ok(image) => {
+      //info!("Image len: {:?}", image.len());
+      //let bytes = BASE64_STANDARD.encode(image.as_bytes());
+      //// [2025-02-18][04:14:15][app_lib::endpoints::image_endpoint][INFO] First bytes: "DAwODAwNDA"
+      //info!("First bytes: {:?}", bytes.split_at(10).0);
+
+      
       let img_data = std::fs::read("temp.png")
         .map_err(|e| format!("Failed to read generated image: {}", e))?;
-        
-      let encoded = BASE64_STANDARD.encode(&img_data);
+
+      let bytes = BASE64_STANDARD.encode(&img_data);
       println!("Generated image encoded successfully");
 
       let _ = std::fs::remove_file("temp.png");
 
-      Ok(encoded)
-    }
-    Err(e) => Err(format!("Failed to generate image: {}", e))
+
+
+
+
+      Ok(bytes)
+    },
+    Err(e) => Err(format!("Failed to generate image: {}", e)),
   }
 }

crates/desktop/src-tauri/src/endpoints/sd.rs

@@ -1,37 +1,27 @@
-#[cfg(feature = "accelerate")]
-extern crate accelerate_src;
-#[cfg(feature = "mkl")]
-extern crate intel_mkl_src;
-
-use std::mem::forget;
-use std::path::Path;
-use std::sync::{Arc, LockResult, RwLock, RwLockReadGuard, RwLockWriteGuard};
-use candle_transformers::models::stable_diffusion;
+use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 
+use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
+use crate::ml::image::save_image_from_tensor::save_image_from_tensor;
+use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
 use crate::ml::infer_clip_text_embeddings::infer_clip_text_embeddings;
-use crate::ml::model_file::StableDiffusionVersion;
-use crate::ml::save_image_from_tensor::save_image_from_tensor;
-use anyhow::{Error as E, Result};
-use candle_core::{DType, Device, IndexOp, Module, Tensor, D};
-use candle_transformers::models::stable_diffusion::vae::DiagonalGaussianDistribution;
-use hf_hub::api::sync::Api;
-use log::info;
-use once_cell::sync::Lazy;
-use rand::Rng;
-use crate::ml::create_inpainting_tensors::create_inpainting_tensors;
-use crate::ml::load_image_file_to_tensor::load_image_file_to_tensor;
-use crate::ml::load_image_file_to_tensor_2::load_image_file_to_tensor_2;
 use crate::ml::model_cache::ModelCache;
-use crate::ml::models::unet_model::UNetModel;
 use crate::ml::prompt_cache::PromptCache;
+use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
 use crate::state::app_config::AppConfig;
+use anyhow::{Error as E, Result};
+use candle_core::{DType, Device, IndexOp, D};
+use candle_transformers::models::stable_diffusion::vae::DiagonalGaussianDistribution;
+use image::DynamicImage;
+use log::info;
+use rand::Rng;
+
 // TODO: Clean up
 
 // Note for Kasisnu: I'm going to start using lifetimes as long as that doesn't slow your velocity
 // Basically the args o this telescopic args struct are guaranteed to live as long as the struct
 // itself with the 'a lifetime.
 pub struct Args<'a> {
-    pub image_path: &'a Path,
+    pub image: &'a DynamicImage,
     pub prompt: String,
     pub uncond_prompt: String,
     pub guidance_scale: Option<f64>,
@@ -40,7 +30,7 @@ pub struct Args<'a> {
     pub prompt_cache: &'a PromptCache,
 }
 
-pub fn run(args: Args<'_>) -> Result<()> {
+pub fn run(args: Args<'_>) -> Result<RgbImage> {
     println!("Starting image generation with the following configuration:");
     println!("  Model: {:?}", args.configs.sd_version);
     println!("  Prompt: {}", args.prompt);
@@ -119,17 +109,14 @@ pub fn run(args: Args<'_>) -> Result<()> {
 
     println!("Loading input image into tensor...");
 
-    let input_image = load_image_file_to_tensor_2(args.image_path)?;
-
-    let input_image = input_image.to_device(&args.configs.device)?
-      .to_dtype(args.configs.dtype)?;
+    let input_image = dynamic_image_to_tensor(
+        args.image, 
+        &args.configs.device,
+        args.configs.dtype)?;
 
     // REFERENCE IMAGE shape: [1, 3, 1024, 1024]
     println!("REFERENCE IMAGE shape: {:?}", input_image.shape());
 
-
-
-
     println!("Calculating start step for diffusion process...");
     let img2img_strength = 0.75f64;
     
@@ -142,17 +129,7 @@ pub fn run(args: Args<'_>) -> Result<()> {
 
     let use_guide_scale = false; // TODO
 
-    let vae_scale = match args.configs.sd_version {
-        StableDiffusionVersion::V1_5
-        | StableDiffusionVersion::V1_5Inpaint
-        | StableDiffusionVersion::V2_1
-        | StableDiffusionVersion::V2Inpaint
-        | StableDiffusionVersion::XlInpaint
-        | StableDiffusionVersion::Xl => 0.18215,
-        StableDiffusionVersion::Turbo => 0.13025,
-    };
-
-
+    let vae_scale = get_vae_scale(args.configs.sd_version);
 
     let encoded_image : DiagonalGaussianDistribution = args.model_cache.vae_encode(&input_image)?;
     let init_latent_dist = encoded_image;
@@ -173,15 +150,6 @@ pub fn run(args: Args<'_>) -> Result<()> {
 
     println!("Latents initialized successfully");
 
-
-
-    let vae_scale = match args.configs.sd_version {
-        StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint | StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Xl => 0.18215,
-        StableDiffusionVersion::Turbo => 0.13025,
-    };
-    
-    
-
     //let encoded_sample = encoded_image.sample()?;
 
     // Encoded sample shape: [1, 4, 128, 128] ... so close. Just needs to be smaller.
@@ -190,8 +158,6 @@ pub fn run(args: Args<'_>) -> Result<()> {
     //let mask_latents = (encoded_sample * vae_scale)?.to_device(&args.model_configs.device)?;
 
     //println!("Mask latents shape: {:?}", mask_latents.shape());
-
-
     
     //let timesteps = scheduler.timesteps().to_vec();
     let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect();
@@ -321,8 +287,12 @@ pub fn run(args: Args<'_>) -> Result<()> {
     let image = (image.clamp(0f32, 1.)? * 255.)?.to_dtype(DType::U8)?;
     println!("Converted to 8-bit format");
 
-    save_image_from_tensor(&image.i(0)?, "temp.png")?;
-    println!("Image generation completed successfully");
+    //save_image_from_tensor(&image.i(0)?, "temp.png")?;
+    //println!("Image generation completed successfully");
+    
+    let image = tensor_to_image_buffer(&image.i(0)?)?;
 
-    Ok(())
+    image.save("temp.png").map_err(E::from)?;
+    
+    Ok(image)
 }

crates/desktop/src-tauri/src/ml/image/dynamic_image_to_tensor.rs

@@ -0,0 +1,17 @@
+use anyhow::Result;
+use candle_core::{DType, Device, Tensor};
+use image::DynamicImage;
+
+pub fn dynamic_image_to_tensor(image: &DynamicImage, device: &Device, datatype: DType) -> Result<Tensor> {
+  let (height, width) = (image.height() as usize, image.width() as usize);
+  let height = height - height % 32;
+  let width = width - width % 32;
+  let image = image.resize_to_fill(width as u32, height as u32, image::imageops::FilterType::CatmullRom);
+  let image = image.to_rgb8();
+  let image = image.into_raw();
+  Ok(Tensor::from_vec(image, (height, width, 3), device)?
+    .permute((2, 0, 1))?
+    .to_dtype(datatype)?
+    .affine(2. / 255., -1.)?
+    .unsqueeze(0)?)
+}

crates/desktop/src-tauri/src/ml/image/load_image_file_to_tensor_2.rs

@@ -1,10 +1,3 @@
-#[cfg(feature = "accelerate")]
-extern crate accelerate_src;
-#[cfg(feature = "mkl")]
-extern crate intel_mkl_src;
-
-use candle_transformers::models::stable_diffusion;
-
 use anyhow::Result;
 use candle_core::{DType, Device, Tensor};
 

crates/desktop/src-tauri/src/ml/image/mod.rs

@@ -0,0 +1,5 @@
+pub mod dynamic_image_to_tensor;
+pub mod load_image_file_to_tensor;
+pub mod load_image_file_to_tensor_2;
+pub mod save_image_from_tensor;
+pub mod tensor_to_image_buffer;

crates/desktop/src-tauri/src/ml/image/save_image_from_tensor.rs

@@ -1,8 +1,3 @@
-#[cfg(feature = "accelerate")]
-extern crate accelerate_src;
-#[cfg(feature = "mkl")]
-extern crate intel_mkl_src;
-
 use anyhow::{Error as E, Result};
 use candle_core::Tensor;
 use image::{ImageBuffer, Rgb};

crates/desktop/src-tauri/src/ml/image/tensor_to_image_buffer.rs

@@ -0,0 +1,18 @@
+use candle_core::Tensor;
+use image::{ImageBuffer, Rgb};
+
+pub type RgbImage = ImageBuffer<Rgb<u8>, Vec<u8>>;
+
+pub fn tensor_to_image_buffer(img: &Tensor) -> anyhow::Result<RgbImage> {
+  let (channel, height, width) = img.dims3()?;
+
+  if channel != 3 {
+    anyhow::bail!("save_image expects an input of shape (3, height, width)")
+  }
+
+  let img = img.permute((1, 2, 0))?.flatten_all()?;
+  let pixels = img.to_vec1::<u8>()?;
+
+  ImageBuffer::from_raw(width as u32, height as u32, pixels)
+    .ok_or_else(|| anyhow::anyhow!("error loading image from raw tensor"))
+}

crates/desktop/src-tauri/src/ml/mod.rs

@@ -1,9 +1,8 @@
 pub mod create_inpainting_tensors;
+pub mod image;
 pub mod infer_clip_text_embeddings;
-pub mod load_image_file_to_tensor;
-pub mod load_image_file_to_tensor_2;
 pub mod model_cache;
 pub mod model_file;
 pub mod models;
 pub mod prompt_cache;
-pub mod save_image_from_tensor;
+pub mod stable_diffusion;

crates/desktop/src-tauri/src/ml/stable_diffusion/get_vae_scale.rs

@@ -0,0 +1,13 @@
+use crate::ml::model_file::StableDiffusionVersion;
+
+pub fn get_vae_scale(sd_version: StableDiffusionVersion) -> f64 {
+  match sd_version {
+    StableDiffusionVersion::V1_5
+    | StableDiffusionVersion::V1_5Inpaint
+    | StableDiffusionVersion::V2_1
+    | StableDiffusionVersion::V2Inpaint
+    | StableDiffusionVersion::XlInpaint
+    | StableDiffusionVersion::Xl => 0.18215,
+    StableDiffusionVersion::Turbo => 0.13025,
+  }
+}

crates/desktop/src-tauri/src/ml/stable_diffusion/mod.rs

@@ -0,0 +1 @@
+pub mod get_vae_scale;

crates/desktop/src-tauri/src/state/app_config.rs

@@ -40,10 +40,9 @@ pub struct AppConfig {
 impl AppConfig {
   
   pub fn init() -> anyhow::Result<Self> {
-    
+    println!("Configuration filename: {}", CONFIG_FILENAME);
     let yaml_configs = load_yaml_configs();
     
-    println!("Configured filename {}", CONFIG_FILENAME);
     println!("Configs: {:?}", yaml_configs);
     
     let device = Device::new_cuda(0)
@@ -83,6 +82,7 @@ fn load_yaml_configs() -> YamlConfig {
     .ok()
     .flatten();
   if let Some(path) = maybe_path {
+    let path = path.canonicalize().unwrap_or(path);
     println!("Attempting to load configs from: {:?}", &path);
     if let Ok(config) = YamlConfig::read_from_file(path) {
       return config;