diff --git a/crates/desktop/src-tauri/src/endpoints/sd.rs b/crates/desktop/src-tauri/src/endpoints/sd.rs index 6b9d513334..54cb724a8d 100644 --- a/crates/desktop/src-tauri/src/endpoints/sd.rs +++ b/crates/desktop/src-tauri/src/endpoints/sd.rs @@ -1,7 +1,6 @@ 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_cache::ModelCache; @@ -9,7 +8,7 @@ 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_core::{DType, IndexOp, D}; use candle_transformers::models::stable_diffusion::vae::DiagonalGaussianDistribution; use image::DynamicImage; use log::info; @@ -25,8 +24,8 @@ pub struct Args<'a> { pub prompt: String, pub uncond_prompt: String, pub guidance_scale: Option, - pub model_cache: &'a ModelCache, pub configs: &'a AppConfig, + pub model_cache: &'a ModelCache, pub prompt_cache: &'a PromptCache, } @@ -45,39 +44,10 @@ pub fn run(args: Args<'_>) -> Result { args.configs.scheduler_steps)?; let seed = args.configs.seed.unwrap_or_else(|| rand::thread_rng().gen()); - println!("Using seed: {}", seed); + println!("Using seed: {}", seed); args.configs.device.set_seed(seed)?; - println!("Initializing Hugging Face API..."); - - let repo = args.configs.sd_version.repo(); - println!("Downloading model files from: {}", repo); - - //println!("Downloading VAE ..."); - //let vae_file = args.api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?; - - //println!("VAE Path: {:?}", &vae_file); - - //println!("Downloading text encoders..."); - //let clip_file = args.api.model(repo.to_string()).get("text_encoder/model.safetensors")?; - //let clip2_file = args.api.model(repo.to_string()).get("text_encoder_2/model.safetensors")?; - //let tokenizer = args.api.model("openai/clip-vit-large-patch14".to_string()).get("tokenizer.json")?; - //let tokenizer2 = args.api.model("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k".to_string()).get("tokenizer.json")?; - - //println!("Building VAE model..."); - //let vae = sd_config.build_vae(vae_file, &device, 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)?; - info!("Checking if prompt is cached"); let maybe_cached = args.prompt_cache.get_copy(&args.prompt)?; @@ -85,7 +55,7 @@ pub fn run(args: Args<'_>) -> Result { info!("Prompt is cached!"); tensor } else { - info!("Prompt is NOT cached!"); + info!("Prompt is NOT cached! Calculating embedding..."); let tensor = infer_clip_text_embeddings( &args.prompt, &args.uncond_prompt, @@ -99,9 +69,7 @@ pub fn run(args: Args<'_>) -> Result { args.configs.dtype, false, // use_guide_scale )?; - args.prompt_cache.store_copy(&args.prompt, &tensor)?; - tensor }; @@ -114,31 +82,27 @@ pub fn run(args: Args<'_>) -> Result { &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; - - let t_start = { - let start = args.configs.scheduler_steps - (args.configs.scheduler_steps as f64 * img2img_strength) as usize; - - println!("Starting from step {} of {} (strength: {})", start, args.configs.scheduler_steps, img2img_strength); - start - }; - - let use_guide_scale = false; // TODO + println!("Reference image shape: {:?}", input_image.shape()); 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; + let init_latent_dist : DiagonalGaussianDistribution = args.model_cache.vae_encode(&input_image)?; println!("Generating latents from input image..."); let latents = (init_latent_dist.sample()? * vae_scale)?.to_device(&args.configs.device)?; println!("Initial latents shape: {:?}", latents.shape()); + println!("Calculating start step for diffusion process..."); + let img2img_strength = 0.75f64; + + let t_start = { + let start = args.configs.scheduler_steps - (args.configs.scheduler_steps as f64 * img2img_strength) as usize; + + println!("Starting from step {} of {} (strength: {})", start, args.configs.scheduler_steps, img2img_strength); + start + }; + let latents = if t_start < scheduler.timesteps().len() { println!("Adding noise to latents..."); let noise = latents.randn_like(0f64, 1f64)?; @@ -150,29 +114,8 @@ pub fn run(args: Args<'_>) -> Result { println!("Latents initialized successfully"); - //let encoded_sample = encoded_image.sample()?; - - // Encoded sample shape: [1, 4, 128, 128] ... so close. Just needs to be smaller. - //println!("Encoded sample shape: {:?}", encoded_sample.shape()); - - //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(); - //// TODO: This shape may need to change for inpainting. - //let latents = Tensor::randn( - // 0f32, - // 1f32, - // (1, 4, args.model_configs.sd_config.height / 8, args.model_configs.sd_config.width / 8), - // &args.model_configs.device, - //)?; - - //let mut latents = (latents * scheduler.init_noise_sigma())?; - //let mut latents = (mask_latents * scheduler.init_noise_sigma())?; - // Initial noise shape: [1, 4, 64, 64] println!("Initial noise shape: {:?}", latents.shape()); @@ -184,66 +127,19 @@ pub fn run(args: Args<'_>) -> Result { if timestep_index < t_start { continue; } - - // println!("Processing step {}/{}", timestep_index + 1, timesteps.len()); - - // let latent_model_input = latents.clone(); - - // // Latent input shape: [1, 4, 64, 64] - // println!("Latent input shape: {:?}", latent_model_input.shape()); - // - // println!("Scaling model input..."); - // let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?; - - // // Scaled input shape: [1, 4, 64, 64] - // println!("Scaled input shape: {:?}", latent_model_input.shape()); - - - // // TODO: THIS IS A HACK - // //let latent_model_input = match HACK_INPAINT_SD_VERSION { - // // StableDiffusionVersion::XlInpaint - // // | StableDiffusionVersion::V2Inpaint - // // | StableDiffusionVersion::V1_5Inpaint => { - // // info!("Concatenating input shape: {:?}", latent_model_input.shape()); - // // info!("Mask shape: {:?}", mask.as_ref().unwrap().shape()); - // // info!("Mask latents shape: {:?}", mask_latents.as_ref().unwrap().shape()); - // // info!("IF THIS FAILS, REVERT THE `HACK_INPAINT_SD_VERSION` 'flag'"); - // // Tensor::cat( - // // &[ - // // &latent_model_input, - // // mask.as_ref().unwrap(), - // // mask_latents.as_ref().unwrap(), - // // ], - // // 1, - // // )? - // // }, - // // _ => latent_model_input, - // //} - // // .to_device(&args.model_configs.device)?; - - - // - // println!("Running UNet inference with timestep {}...", timestep); - - // let noise_pred = args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings)?; - - // 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!("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 args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings) { Ok(pred) => { println!("UNet inference successful"); @@ -261,34 +157,22 @@ pub fn run(args: Args<'_>) -> Result { println!("Step {}/{} completed", timestep_index + 1, timesteps.len()); } - //println!("Diffusion process completed, decoding image..."); - //let image = args.model_cache.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_from_tensor(&image.i(0)?, "temp.png")?; - //println!("Image generation completed successfully"); - - - println!("Diffusion process completed, decoding image..."); let image = args.model_cache.vae_decode(&(latents / vae_scale)?)?; - //let image = args.model_cache.vae_decode(&(latents.div(vae_scale)?))?; + println!("VAE decode completed"); - //let image = ((image.div(2.))? + 0.5)?.to_device(&Device::Cpu)?; - 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_from_tensor(&image.i(0)?, "temp.png")?; - //println!("Image generation completed successfully"); + println!("Scaling image back"); + let image = ((image / 2.)? + 0.5)?; + + // TODO(bt,2025-08-18): This normalization is slow. + println!("Normalized image values"); + let image = (image.clamp(0f32, 1.)? * 255.)?; + + println!("Convert to int8"); + let image = image.to_dtype(DType::U8)?; + + println!("Converted to 8-bit format"); let image = tensor_to_image_buffer(&image.i(0)?)?;