diff --git a/.idea/dictionaries/bt.xml b/.idea/dictionaries/bt.xml index d253914687..75f4f76930 100644 --- a/.idea/dictionaries/bt.xml +++ b/.idea/dictionaries/bt.xml @@ -47,6 +47,7 @@ techcrunch transcoded transcodings + unet unpermitted upserting upserts diff --git a/crates/desktop/src-tauri/src/model/gpu_detector.rs b/crates/desktop/src-tauri/src/model/gpu_detector.rs new file mode 100644 index 0000000000..e69de29bb2 diff --git a/crates/desktop/src-tauri/src/model/mod.rs b/crates/desktop/src-tauri/src/model/mod.rs index 59df3c0797..ce73b7aaf7 100644 --- a/crates/desktop/src-tauri/src/model/mod.rs +++ b/crates/desktop/src-tauri/src/model/mod.rs @@ -1,3 +1,6 @@ -pub mod save_image_from_tensor; +pub mod gpu_detector; pub mod infer_clip_text_embeddings; -pub mod model_file; \ No newline at end of file +pub mod model_file; +pub mod save_image_from_tensor; +pub mod unet_model; +pub mod model_cache; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/model/model_cache.rs b/crates/desktop/src-tauri/src/model/model_cache.rs new file mode 100644 index 0000000000..0857ddfc2b --- /dev/null +++ b/crates/desktop/src-tauri/src/model/model_cache.rs @@ -0,0 +1,21 @@ +use std::sync::{Arc, RwLock}; +use once_cell::sync::Lazy; +use crate::model::unet_model::UNetModel; + + +static MODEL_CACHE: Lazy>> = Lazy::new(|| Arc::new(RwLock::new(ModelCache::new()))); + + +// Simple registry for now. We can build complex machinery that aids in +// VRAM utilization, disk space saving, intelligent scheduling, etc. in the future. +pub struct ModelCache { + unet: Option, +} + +impl ModelCache { + fn new() -> Self { + Self { + unet: None + } + } +} diff --git a/crates/desktop/src-tauri/src/model/unet_model.rs b/crates/desktop/src-tauri/src/model/unet_model.rs new file mode 100644 index 0000000000..a1eee68c3c --- /dev/null +++ b/crates/desktop/src-tauri/src/model/unet_model.rs @@ -0,0 +1,72 @@ +use anyhow::anyhow; +use candle_core::{DType, Device, Tensor}; +use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel; +use candle_transformers::models::stable_diffusion::StableDiffusionConfig; +use std::path::Path; +use std::sync::{Arc, RwLock}; +use log::debug; + +#[derive(Clone)] +pub struct UNetModel { + // Kasisnu - a few notes: + // + // (1) This struct has what Rust terms as "interior mutability". We can now pass around + // "immutable" references to the struct, yet mutate what's inside with non-mut (!) functions. + // The callers don't know, and because the struct has safety guarantees of "Send + Sync", it's + // provably sound. This allows multiple threads / scopes to mutate the structure. They just wait + // on the lock/mutex to synchronize. + // + // (2) Another cool thing we can do is mark this struct as "Clone" since the Arc pointer is + // super cheap to duplicate (you'd never do that with a wide type or anything beyond int64 + // size). It makes passing this structure along smooth as butter. The only penalty we pay is the + // unlocking of the mutex (actually rwlock in this case) and the (minor) indirection behind Arc. + // We don't mark it Copy, though, since it does ref counting and manages its own allocation, which + // shouldn't be implicit. + model: Arc>, +} + +impl UNetModel { + pub fn new>( + sd_config: &StableDiffusionConfig, + unet_weights_file: P, + device: &Device, + dtype: DType, + ) -> anyhow::Result { + debug!("building unet model..."); + + let unet = sd_config.build_unet( + unet_weights_file, + &device, + 4, + false, + dtype + )?; + + Ok(Self { + model: Arc::new(RwLock::new(unet)), + }) + } + + pub fn inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result { + match self.model.read() { + Err(err) => return Err(anyhow!("lock error: {:?}", err)), + Ok(model) => { + debug!("Running model prediction..."); + let inference_result = model.forward(&latent_model_input, timestep, &text_embeddings); + + let noise_pred = match inference_result { + Ok(pred) => { + debug!("Inference successful"); + pred + }, + Err(e) => { + debug!("UNet inference failed with error: {}", e); + return Err(anyhow!("UNet inference failed: {}", e)); + } + }; + + Ok(noise_pred) + } + } + } +} diff --git a/crates/desktop/src-tauri/src/sd.rs b/crates/desktop/src-tauri/src/sd.rs index 5970df5315..3070c4ba7e 100644 --- a/crates/desktop/src-tauri/src/sd.rs +++ b/crates/desktop/src-tauri/src/sd.rs @@ -11,6 +11,7 @@ use crate::model::save_image_from_tensor::save_image_from_tensor; use anyhow::{Error as E, Result}; use candle_core::{DType, Device, IndexOp, Module, Tensor, D}; use rand::Rng; +use crate::model::unet_model::UNetModel; pub struct Args { pub prompt: String, @@ -65,12 +66,14 @@ pub fn run(args: Args) -> Result<()> { 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); @@ -78,15 +81,21 @@ pub fn run(args: Args) -> Result<()> { 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!("VAE Path: {:?}", &vae_file); + println!("UNET Path: {:?}", &unet_file); + 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..."); + println!("Building VAE model..."); let vae = sd_config.build_vae(vae_file, &device, dtype)?; - let unet = sd_config.build_unet(unet_file, &device, 4, false, dtype)?; + + println!("Building UNET model..."); + //let unet = sd_config.build_unet(unet_file, &device, 4, false, dtype)?; + let unet = UNetModel::new(&sd_config, unet_file, &device, dtype)?; // Build text encoders println!("Building text encoders..."); @@ -152,16 +161,20 @@ pub fn run(args: Args) -> Result<()> { 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)); - } - }; + + let noise_pred = unet.inference(&latent_model_input, timestep as f64, &text_embeddings)?; + + //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..."); diff --git a/script/launch_realtime_desktop.sh b/script/launch_realtime_desktop.sh index 55f0faec6a..15ba9f4ed2 100755 --- a/script/launch_realtime_desktop.sh +++ b/script/launch_realtime_desktop.sh @@ -8,4 +8,4 @@ root_dir=$(pwd) export TAURI_FRONTEND_PATH="${root_dir}/frontend/realtime" -cargo tauri dev +cargo tauri dev --no-watch