a hacky version

crates/desktop/src-tauri/Cargo.toml

@@ -26,3 +26,14 @@ tauri-plugin-log = "2.0.0-rc"
 once_cell = "1.20.3"
 image = "0.25.5"
 base64 = "0.22.1"
+candle-core = { version = "0.8.2", features = ["cuda"] }
+candle-transformers = { version = "0.8.2", features = ["cuda"] }
+candle-nn = { version = "0.8.2", features = ["cuda"] }
+tokenizers = "0.15.2"
+hf-hub = "0.3.2"
+anyhow = "1.0"
+rand = "0.8"
+
+[features]
+default = ["cuda"]
+cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]

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

@@ -1,34 +1,51 @@
+use anyhow::Result;
+use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
 use std::io::Cursor;
-use base64::Engine;
-use base64::prelude::{BASE64_STANDARD, BASE64_URL_SAFE};
-use image::codecs::png::{CompressionType, FilterType, PngEncoder};
-use image::ImageReader;
+use crate::sd::{self, StableDiffusionVersion, Args};
+use candle_core::Device;
+
+const PROMPT: &str = "A beautiful landscape with mountains and a lake";
 
 #[tauri::command]
 pub fn infer_image(image: &str) -> Result<String, String> {
-  println!("infer_image called; processing image...");
+    println!("infer_image called; generating image with SDXL Turbo...");
+    
+    // Check CUDA availability
+    let use_cpu = match Device::new_cuda(0) {
+        Ok(_) => {
+            println!("CUDA device available");
+            false
+        }
+        Err(e) => {
+            println!("CUDA not available ({}), falling back to CPU", e);
+            true
+        }
+    };
+    
+    let args = Args {
+        prompt: PROMPT.to_string(),
+        uncond_prompt: "".to_string(),
+        cpu: use_cpu,
+        height: Some(512),
+        width: Some(512),
+        n_steps: Some(1),
+        num_samples: 1,
+        seed: None,
+        sd_version: StableDiffusionVersion::Turbo,
+        guidance_scale: Some(0.0),
+    };
 
-  let bytes = BASE64_STANDARD.decode(image)
-    .map_err(|err| format!("Base64 decode error: {}", err))?;
-
-  let image = ImageReader::new(Cursor::new(bytes))
-    .with_guessed_format()
-    .map_err(|err| format!("Image format error: {}", err))?
-    .decode()
-    .map_err(|err| format!("Image decode error: {}", err))?;
-
-  let rotated_image = image.rotate180();
-  
-  let mut buffer = Vec::new(); // TODO: Preallocate, better data structure
-  let mut writer = Cursor::new(&mut buffer);
-  
-  let encoder = PngEncoder::new_with_quality(&mut writer, CompressionType::Fast, FilterType::Adaptive);
-  rotated_image.write_with_encoder(encoder)
-    .map_err(|err| format!("Image encoding error: {}", err))?;
-
-  let encoded = BASE64_STANDARD.encode(&buffer);
-  
-  println!("Encoded: {:?}", encoded.split_at(10).0);
-  
-  Ok(encoded)
+    match sd::run(args) {
+        Ok(_) => {
+            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);
+            println!("Generated image encoded successfully");
+            
+            let _ = std::fs::remove_file("temp.png");
+            
+            Ok(encoded)
+        }
+        Err(e) => Err(format!("Failed to generate image: {}", e))
+    }
 }

crates/desktop/src-tauri/src/lib.rs

@@ -1,4 +1,5 @@
 mod image_endpoint;
+mod sd;
 
 use std::sync::{Arc, RwLock};
 use once_cell::sync::Lazy;

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

@@ -0,0 +1,385 @@
+#[cfg(feature = "accelerate")]
+extern crate accelerate_src;
+
+#[cfg(feature = "mkl")]
+extern crate intel_mkl_src;
+
+use candle_transformers::models::stable_diffusion;
+use std::ops::Div;
+
+use anyhow::{Error as E, Result};
+use candle_core::{DType, Device, IndexOp, Tensor, D, Module};
+use image::{ImageBuffer, Rgb};
+use rand::Rng;
+use stable_diffusion::vae::AutoEncoderKL;
+use tokenizers::Tokenizer;
+
+// Helper function to save image
+fn save_image<P: AsRef<std::path::Path>>(img: &Tensor, p: P) -> Result<()> {
+    println!("Saving image to {:?}", p.as_ref());
+    let p = p.as_ref();
+    let (channel, height, width) = img.dims3()?;
+    println!("Image dimensions: {}x{} with {} channels", width, height, channel);
+    
+    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>()?;
+    println!("Converting tensor to image buffer...");
+    
+    let image: ImageBuffer<Rgb<u8>, Vec<u8>> =
+        match ImageBuffer::from_raw(width as u32, height as u32, pixels) {
+            Some(image) => image,
+            None => anyhow::bail!("error saving image {p:?}"),
+        };
+    println!("Successfully created image buffer");
+    
+    image.save(p).map_err(E::from)?;
+    println!("Successfully saved image to disk");
+    Ok(())
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum StableDiffusionVersion {
+    V1_5,
+    V1_5Inpaint,
+    V2_1,
+    V2Inpaint,
+    Xl,
+    XlInpaint,
+    Turbo,
+}
+
+pub struct Args {
+    pub prompt: String,
+    pub uncond_prompt: String,
+    pub cpu: bool,
+    pub height: Option<usize>,
+    pub width: Option<usize>,
+    pub n_steps: Option<usize>,
+    pub num_samples: usize,
+    pub seed: Option<u64>,
+    pub sd_version: StableDiffusionVersion,
+    pub guidance_scale: Option<f64>,
+}
+
+impl StableDiffusionVersion {
+    fn repo(&self) -> &'static str {
+        match self {
+            Self::XlInpaint => "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+            Self::Xl => "stabilityai/stable-diffusion-xl-base-1.0",
+            Self::V2Inpaint => "stabilityai/stable-diffusion-2-inpainting",
+            Self::V2_1 => "stabilityai/stable-diffusion-2-1",
+            Self::V1_5 => "runwayml/stable-diffusion-v1-5",
+            Self::V1_5Inpaint => "stable-diffusion-v1-5/stable-diffusion-inpainting",
+            Self::Turbo => "stabilityai/sdxl-turbo",
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum ModelFile {
+    Tokenizer,
+    Tokenizer2,
+    Clip,
+    Clip2,
+    Unet,
+    Vae,
+}
+
+impl ModelFile {
+    fn get(
+        &self,
+        filename: Option<String>,
+        version: StableDiffusionVersion,
+        use_f16: bool,
+    ) -> Result<std::path::PathBuf> {
+        use hf_hub::api::sync::Api;
+        match filename {
+            Some(filename) => Ok(std::path::PathBuf::from(filename)),
+            None => {
+                let (repo, path) = match self {
+                    Self::Tokenizer => {
+                        let tokenizer_repo = match version {
+                            StableDiffusionVersion::V1_5
+                            | StableDiffusionVersion::V2_1
+                            | StableDiffusionVersion::V1_5Inpaint
+                            | StableDiffusionVersion::V2Inpaint => "openai/clip-vit-base-patch32",
+                            StableDiffusionVersion::Xl
+                            | StableDiffusionVersion::XlInpaint
+                            | StableDiffusionVersion::Turbo => {
+                                "openai/clip-vit-large-patch14"
+                            }
+                        };
+                        (tokenizer_repo, "tokenizer.json")
+                    }
+                    Self::Tokenizer2 => {
+                        ("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", "tokenizer.json")
+                    }
+                    Self::Clip => (version.repo(), "text_encoder/model.safetensors"),
+                    Self::Clip2 => (version.repo(), "text_encoder_2/model.safetensors"),
+                    Self::Unet => (version.repo(), "unet/diffusion_pytorch_model.safetensors"),
+                    Self::Vae => (version.repo(), "vae/diffusion_pytorch_model.safetensors"),
+                };
+                let filename = Api::new()?.model(repo.to_string()).get(path)?;
+                Ok(filename)
+            }
+        }
+    }
+}
+
+pub fn run(args: Args) -> Result<()> {
+    println!("Starting image generation with the following configuration:");
+    println!("  Model: {:?}", args.sd_version);
+    println!("  Prompt: {}", args.prompt);
+    println!("  Steps: {}", args.n_steps.unwrap_or(1));
+    println!("  Using CPU: {}", args.cpu);
+    
+    let device = if args.cpu { 
+        println!("Using CPU for computation");
+        Device::Cpu 
+    } else { 
+        println!("Attempting to use CUDA device");
+        match Device::new_cuda(0) {
+            Ok(cuda_device) => {
+                println!("Successfully initialized CUDA device");
+                cuda_device
+            }
+            Err(e) => {
+                println!("Failed to initialize CUDA device: {}. Falling back to CPU", e);
+                Device::Cpu
+            }
+        }
+    };
+    let dtype = DType::F32;
+    println!("Using {:?} data type", dtype);
+
+    let sd_config = match args.sd_version {
+        StableDiffusionVersion::Turbo => {
+            println!("Configuring SDXL Turbo");
+            stable_diffusion::StableDiffusionConfig::sdxl_turbo(None, args.height, args.width)
+        }
+        _ => {
+            println!("Configuring SD v2.1");
+            stable_diffusion::StableDiffusionConfig::v2_1(None, args.height, args.width)
+        }
+    };
+    println!("Model dimensions: {}x{}", sd_config.width, sd_config.height);
+
+    println!("Building scheduler...");
+    let mut scheduler = sd_config.build_scheduler(args.n_steps.unwrap_or(1))?;
+    let seed = args.seed.unwrap_or_else(|| rand::thread_rng().gen());
+    println!("Using seed: {}", seed);
+    device.set_seed(seed)?;
+
+    println!("Initializing Hugging Face API...");
+    let api = hf_hub::api::sync::Api::new()?;
+    let repo = args.sd_version.repo();
+    println!("Downloading model files from: {}", repo);
+    
+    println!("Downloading VAE and UNet...");
+    let vae_file = api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?;
+    let unet_file = api.model(repo.to_string()).get("unet/diffusion_pytorch_model.safetensors")?;
+    
+    println!("Downloading text encoders...");
+    let clip_file = api.model(repo.to_string()).get("text_encoder/model.safetensors")?;
+    let clip2_file = api.model(repo.to_string()).get("text_encoder_2/model.safetensors")?;
+    let tokenizer = api.model("openai/clip-vit-large-patch14".to_string()).get("tokenizer.json")?;
+    let tokenizer2 = api.model("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k".to_string()).get("tokenizer.json")?;
+
+    println!("Building models...");
+    let vae = sd_config.build_vae(vae_file, &device, dtype)?;
+    let unet = sd_config.build_unet(unet_file, &device, 4, false, dtype)?;
+    
+    // Build text encoders
+    println!("Building text encoders...");
+    let text_model = stable_diffusion::build_clip_transformer(&sd_config.clip, clip_file, &device, dtype)?;
+    let text_model2 = stable_diffusion::build_clip_transformer(sd_config.clip2.as_ref().unwrap(), clip2_file, &device, dtype)?;
+    
+    // Tokenize and encode text
+    println!("Processing text prompt...");
+    let tokenizer = tokenizers::Tokenizer::from_file(tokenizer).map_err(E::msg)?;
+    let tokenizer2 = tokenizers::Tokenizer::from_file(tokenizer2).map_err(E::msg)?;
+    
+    println!("Preparing text embeddings...");
+    let which = match args.sd_version {
+        StableDiffusionVersion::Xl
+        | StableDiffusionVersion::XlInpaint
+        | StableDiffusionVersion::Turbo => vec![true, false],
+        _ => vec![true],
+    };
+    
+    let text_embeddings = which
+        .iter()
+        .map(|first| {
+            text_embeddings(
+                &args.prompt,
+                &args.uncond_prompt,
+                None, // tokenizer
+                None, // clip_weights
+                None, // clip2_weights
+                args.sd_version,
+                &sd_config,
+                false, // use_f16
+                &device,
+                dtype,
+                false, // use_guide_scale
+                *first,
+            )
+        })
+        .collect::<Result<Vec<_>>>()?;
+
+    let text_embeddings = Tensor::cat(&text_embeddings, D::Minus1)?;
+    println!("Text embeddings shape: {:?}", text_embeddings.shape());
+
+    println!("Generating initial noise...");
+    let timesteps = scheduler.timesteps().to_vec();
+    let latents = Tensor::randn(
+        0f32,
+        1f32,
+        (1, 4, sd_config.height / 8, sd_config.width / 8),
+        &device,
+    )?;
+    let latents = (latents * scheduler.init_noise_sigma())?;
+    println!("Initial noise shape: {:?}", latents.shape());
+
+    println!("Starting diffusion process...");
+    let mut latents = latents;
+    for (timestep_index, &timestep) in timesteps.iter().enumerate() {
+        println!("Processing step {}/{}", timestep_index + 1, timesteps.len());
+        let latent_model_input = latents.clone();
+        println!("Latent input shape: {:?}", latent_model_input.shape());
+        
+        println!("Scaling model input...");
+        let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?;
+        println!("Scaled input shape: {:?}", latent_model_input.shape());
+        
+        println!("Running UNet inference with timestep {}...", timestep);
+        let noise_pred = match unet.forward(&latent_model_input, timestep as f64, &text_embeddings) {
+            Ok(pred) => {
+                println!("UNet inference successful");
+                pred
+            }
+            Err(e) => {
+                println!("UNet inference failed with error: {}", e);
+                return Err(anyhow::anyhow!("UNet inference failed: {}", e));
+            }
+        };
+        println!("Noise prediction shape: {:?}", noise_pred.shape());
+        
+        println!("Applying scheduler step...");
+        latents = scheduler.step(&noise_pred, timestep, &latents)?;
+        println!("Step {}/{} completed", timestep_index + 1, timesteps.len());
+    }
+
+    println!("Diffusion process completed, decoding image...");
+    let image = vae.decode(&(latents / 0.13025)?)?;
+    println!("VAE decode completed");
+    let image = ((image / 2.)? + 0.5)?.to_device(&Device::Cpu)?;
+    println!("Normalized image values");
+    let image = (image.clamp(0f32, 1.)? * 255.)?.to_dtype(DType::U8)?;
+    println!("Converted to 8-bit format");
+    
+    save_image(&image.i(0)?, "temp.png")?;
+    println!("Image generation completed successfully");
+
+    Ok(())
+}
+
+#[allow(clippy::too_many_arguments)]
+fn text_embeddings(
+    prompt: &str,
+    uncond_prompt: &str,
+    tokenizer: Option<String>,
+    clip_weights: Option<String>,
+    clip2_weights: Option<String>,
+    sd_version: StableDiffusionVersion,
+    sd_config: &stable_diffusion::StableDiffusionConfig,
+    use_f16: bool,
+    device: &Device,
+    dtype: DType,
+    use_guide_scale: bool,
+    first: bool,
+) -> Result<Tensor> {
+    let tokenizer_file = if first {
+        ModelFile::Tokenizer
+    } else {
+        ModelFile::Tokenizer2
+    };
+    let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
+    let tokenizer = Tokenizer::from_file(tokenizer).map_err(E::msg)?;
+    let pad_id = match &sd_config.clip.pad_with {
+        Some(padding) => *tokenizer.get_vocab(true).get(padding.as_str()).unwrap(),
+        None => *tokenizer.get_vocab(true).get("<|endoftext|>").unwrap(),
+    };
+    
+    println!("Running with prompt \"{prompt}\".");
+    let mut tokens = tokenizer
+        .encode(prompt, true)
+        .map_err(E::msg)?
+        .get_ids()
+        .to_vec();
+    if tokens.len() > sd_config.clip.max_position_embeddings {
+        anyhow::bail!(
+            "the prompt is too long, {} > max-tokens ({})",
+            tokens.len(),
+            sd_config.clip.max_position_embeddings
+        )
+    }
+    while tokens.len() < sd_config.clip.max_position_embeddings {
+        tokens.push(pad_id)
+    }
+    let tokens = Tensor::new(tokens.as_slice(), device)?.unsqueeze(0)?;
+
+    println!("Building the Clip transformer.");
+    let clip_weights_file = if first {
+        ModelFile::Clip
+    } else {
+        ModelFile::Clip2
+    };
+    let clip_weights = if first {
+        clip_weights_file.get(clip_weights, sd_version, use_f16)?
+    } else {
+        clip_weights_file.get(clip2_weights, sd_version, use_f16)?
+    };
+    let clip_config = if first {
+        &sd_config.clip
+    } else {
+        sd_config.clip2.as_ref().unwrap()
+    };
+    let text_model = stable_diffusion::build_clip_transformer(clip_config, clip_weights, device, DType::F32)?;
+    
+    let text_embeddings = text_model.forward(&tokens)?;
+
+    let text_embeddings = if use_guide_scale {
+        let mut uncond_tokens = tokenizer
+            .encode(uncond_prompt, true)
+            .map_err(E::msg)?
+            .get_ids()
+            .to_vec();
+        if uncond_tokens.len() > sd_config.clip.max_position_embeddings {
+            anyhow::bail!(
+                "the negative prompt is too long, {} > max-tokens ({})",
+                uncond_tokens.len(),
+                sd_config.clip.max_position_embeddings
+            )
+        }
+        while uncond_tokens.len() < sd_config.clip.max_position_embeddings {
+            uncond_tokens.push(pad_id)
+        }
+
+        let uncond_tokens = Tensor::new(uncond_tokens.as_slice(), device)?.unsqueeze(0)?;
+        let uncond_embeddings = text_model.forward(&uncond_tokens)?;
+
+        Tensor::cat(&[uncond_embeddings, text_embeddings], 0)?.to_dtype(dtype)?
+    } else {
+        text_embeddings.to_dtype(dtype)?
+    };
+    Ok(text_embeddings)
+}
+
+// fn main() -> Result<()> {
+//     let args = Args::parse();
+//     run(args)
+// }