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