From 5faffc014a0ba6923d81251f64561a89ff0193e2 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 14 Feb 2025 17:24:26 -0500 Subject: [PATCH] better memory organization (wip) --- .idea/dictionaries/bt.xml | 1 + .../desktop/src-tauri/src/image_endpoint.rs | 128 ++++++++++++++---- crates/desktop/src-tauri/src/lib.rs | 1 + crates/desktop/src-tauri/src/ml/mod.rs | 2 + .../desktop/src-tauri/src/ml/model_cache.rs | 90 ++++++++++++ crates/desktop/src-tauri/src/ml/models/mod.rs | 2 + .../src-tauri/src/ml/models/unet_model.rs | 35 +++++ .../models/unet_model_cached.rs} | 2 +- crates/desktop/src-tauri/src/model/mod.rs | 2 - .../src-tauri/src/model/model_cache.rs | 21 --- crates/desktop/src-tauri/src/sd.rs | 108 +++++++-------- 11 files changed, 286 insertions(+), 106 deletions(-) create mode 100644 crates/desktop/src-tauri/src/ml/mod.rs create mode 100644 crates/desktop/src-tauri/src/ml/model_cache.rs create mode 100644 crates/desktop/src-tauri/src/ml/models/mod.rs create mode 100644 crates/desktop/src-tauri/src/ml/models/unet_model.rs rename crates/desktop/src-tauri/src/{model/unet_model.rs => ml/models/unet_model_cached.rs} (98%) delete mode 100644 crates/desktop/src-tauri/src/model/model_cache.rs diff --git a/.idea/dictionaries/bt.xml b/.idea/dictionaries/bt.xml index 75f4f76930..489e55b875 100644 --- a/.idea/dictionaries/bt.xml +++ b/.idea/dictionaries/bt.xml @@ -31,6 +31,7 @@ hifigan hparams huggingface + kasisnu lipsync mocap mochi diff --git a/crates/desktop/src-tauri/src/image_endpoint.rs b/crates/desktop/src-tauri/src/image_endpoint.rs index 4ceb6c4153..bd249aee18 100644 --- a/crates/desktop/src-tauri/src/image_endpoint.rs +++ b/crates/desktop/src-tauri/src/image_endpoint.rs @@ -1,11 +1,22 @@ +use std::env::args; +use std::sync::{Arc, RwLock}; use crate::model::model_file::StableDiffusionVersion; use crate::sd::{self, Args}; use anyhow::Result; use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; -use candle_core::Device; +use candle_core::{DType, Device}; +use candle_transformers::models::stable_diffusion; +use once_cell::sync::Lazy; +use crate::ml::model_cache::ModelCache; +use crate::ml::models::unet_model::UNetModel; const PROMPT: &str = "A beautiful landscape with mountains and a lake"; +const HEIGHT: usize = 512; +const WIDTH: usize = 512; + +static MODEL_CACHE: Lazy>>> = Lazy::new(|| Arc::new(RwLock::new(None))); + #[tauri::command] pub fn infer_image(image: &str) -> Result { println!("infer_image called; generating image with SDXL Turbo..."); @@ -22,30 +33,101 @@ pub fn infer_image(image: &str) -> Result { } }; - 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 sd_version = StableDiffusionVersion::Turbo; + + let sd_config = match sd_version { + StableDiffusionVersion::Turbo => { + println!("Configuring SDXL Turbo"); + stable_diffusion::StableDiffusionConfig::sdxl_turbo(None, Some(HEIGHT), Some(WIDTH)) + } + _ => { + println!("Configuring SD v2.1"); + stable_diffusion::StableDiffusionConfig::v2_1(None, Some(HEIGHT), Some(WIDTH)) + } }; + + let api = hf_hub::api::sync::Api::new() + .map_err(|err| err.to_string())?; + + + match MODEL_CACHE.write() { + Err(err) => return Err(err.to_string()), + Ok(mut maybe_cache) => { + match &*maybe_cache { + None => { + let device = if use_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 cache = ModelCache::new(device, sd_version.clone(), sd_config.clone()) + .map_err(|err| err.to_string())?; + + *maybe_cache = Some(cache); + } + Some(_model) => { + println!("Model cache already exists"); + } + } - 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)) } + + + + match MODEL_CACHE.read() { + Err(err) => return Err(err.to_string()), + Ok(maybe_model) => { + match &*maybe_model { + None => { + Err("model does not exist".to_string()) + } + Some(model_cache) => { + + let args = Args { + api, + prompt: PROMPT.to_string(), + uncond_prompt: "".to_string(), + cpu: use_cpu, + height: Some(HEIGHT), + width: Some(WIDTH), + n_steps: Some(1), + num_samples: 1, + seed: None, + sd_version, + guidance_scale: Some(0.0), + model_cache, + }; + + 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)) + } + + } + } + } + } + } diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index 8a2645a2da..e8d63d2419 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -2,6 +2,7 @@ mod flip_image; mod image_endpoint; mod sd; pub mod model; +pub mod ml; use std::sync::{Arc, RwLock}; use once_cell::sync::Lazy; diff --git a/crates/desktop/src-tauri/src/ml/mod.rs b/crates/desktop/src-tauri/src/ml/mod.rs new file mode 100644 index 0000000000..e8a187da65 --- /dev/null +++ b/crates/desktop/src-tauri/src/ml/mod.rs @@ -0,0 +1,2 @@ +pub mod model_cache; +pub mod models; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/ml/model_cache.rs b/crates/desktop/src-tauri/src/ml/model_cache.rs new file mode 100644 index 0000000000..1732e3714f --- /dev/null +++ b/crates/desktop/src-tauri/src/ml/model_cache.rs @@ -0,0 +1,90 @@ +use crate::ml::models::unet_model::UNetModel; +use crate::model::model_file::StableDiffusionVersion; +use anyhow::anyhow; +use candle_core::{DType, Device, Tensor}; +use candle_transformers::models::stable_diffusion::StableDiffusionConfig; +use hf_hub::api::sync::Api; +use std::sync::{Arc, RwLock}; + +// Simple registry for now. We can build complex machinery that aids in +// VRAM utilization, disk space saving, intelligent scheduling, etc. in the future. +#[derive(Clone)] +pub struct ModelCache { + // TODO(bt,2025-02-14): Support for more than one device in the future. + device: Device, + hf_api: Api, + sd_config: StableDiffusionConfig, + sd_version: StableDiffusionVersion, + unet: Arc>>, +} + +impl ModelCache { + pub fn new(device: Device, sd_version: StableDiffusionVersion, sd_config: StableDiffusionConfig) -> anyhow::Result { + Ok(Self { + device, + hf_api: Api::new()?, + unet: Arc::new(RwLock::new(None)), + sd_config, + sd_version, + }) + } + + pub fn unet_inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result { + println!("unet_inference"); + match self.try_unet_inference(latent_model_input, timestep, text_embeddings) { + Err(err) => Err(err), + Ok(Some(result)) => Ok(result), + Ok(None) => { + // Model wasn't previously loaded + self.load_unet()?; + self.try_unet_inference(latent_model_input, timestep, text_embeddings)? + .ok_or(anyhow!("model was not previously loaded")) + } + } + } + + fn load_unet(&self) -> anyhow::Result<()> { + println!("load_unet"); + match self.unet.write() { + Err(err) => return Err(anyhow!("lock error: {:?}", err)), + Ok(mut maybe_model) => { + match &*maybe_model { + Some(_model) => {} // Fall through + None => { + let repo = self.sd_version.repo(); + + println!("Downloading UNET model files from: {} ...", repo); + + let unet_file = self.hf_api.model(repo.to_string()) + .get("unet/diffusion_pytorch_model.safetensors") + .map_err(|err| anyhow!("error fetching model: {:?}", err))?; + + println!("Building UNET model... (3)"); + let unet = UNetModel::new(&self.sd_config, unet_file, &self.device, DType::F32) + .map_err(|err| anyhow!("error initializing unet model: {:?}", err))?; + + println!("Built UNET..."); + *maybe_model = Some(unet); + } + } + } + } + Ok(()) + } + + pub fn try_unet_inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result> { + println!("try_unet_inference"); + match self.unet.read() { + Ok(mut model) => { + match &*model { + None => Ok(None), // Model is not yet loaded + Some(model) => { + Ok(Some(model.inference(latent_model_input, timestep, text_embeddings)?)) + } + } + } + Err(err) => Err(anyhow!("lock error: {:?}", err)), + } + } + +} diff --git a/crates/desktop/src-tauri/src/ml/models/mod.rs b/crates/desktop/src-tauri/src/ml/models/mod.rs new file mode 100644 index 0000000000..8631dad75c --- /dev/null +++ b/crates/desktop/src-tauri/src/ml/models/mod.rs @@ -0,0 +1,2 @@ +pub mod unet_model; +pub mod unet_model_cached; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/ml/models/unet_model.rs b/crates/desktop/src-tauri/src/ml/models/unet_model.rs new file mode 100644 index 0000000000..2483525c96 --- /dev/null +++ b/crates/desktop/src-tauri/src/ml/models/unet_model.rs @@ -0,0 +1,35 @@ +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; + +pub struct UNetModel { + model: UNet2DConditionModel, +} + +impl UNetModel { + pub fn new>( + sd_config: &StableDiffusionConfig, + unet_weights_file: P, + device: &Device, + dtype: DType, + ) -> anyhow::Result { + println!("building unet model... (1)"); + + let model = sd_config.build_unet( + unet_weights_file, + &device, + 4, + false, + dtype + )?; + + Ok(Self { + model, + }) + } + + pub fn inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result { + Ok(self.model.forward(&latent_model_input, timestep, &text_embeddings)?) + } +} diff --git a/crates/desktop/src-tauri/src/model/unet_model.rs b/crates/desktop/src-tauri/src/ml/models/unet_model_cached.rs similarity index 98% rename from crates/desktop/src-tauri/src/model/unet_model.rs rename to crates/desktop/src-tauri/src/ml/models/unet_model_cached.rs index a1eee68c3c..eeb136da61 100644 --- a/crates/desktop/src-tauri/src/model/unet_model.rs +++ b/crates/desktop/src-tauri/src/ml/models/unet_model_cached.rs @@ -32,7 +32,7 @@ impl UNetModel { device: &Device, dtype: DType, ) -> anyhow::Result { - debug!("building unet model..."); + println!("building unet model... (2)"); let unet = sd_config.build_unet( unet_weights_file, diff --git a/crates/desktop/src-tauri/src/model/mod.rs b/crates/desktop/src-tauri/src/model/mod.rs index ce73b7aaf7..ebfcd678cd 100644 --- a/crates/desktop/src-tauri/src/model/mod.rs +++ b/crates/desktop/src-tauri/src/model/mod.rs @@ -2,5 +2,3 @@ pub mod gpu_detector; pub mod infer_clip_text_embeddings; 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 deleted file mode 100644 index 0857ddfc2b..0000000000 --- a/crates/desktop/src-tauri/src/model/model_cache.rs +++ /dev/null @@ -1,21 +0,0 @@ -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/sd.rs b/crates/desktop/src-tauri/src/sd.rs index 5f6811284c..d8d174b23b 100644 --- a/crates/desktop/src-tauri/src/sd.rs +++ b/crates/desktop/src-tauri/src/sd.rs @@ -11,15 +11,21 @@ use crate::model::model_file::StableDiffusionVersion; 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 hf_hub::api::sync::Api; use once_cell::sync::Lazy; use rand::Rng; -use crate::model::model_cache::ModelCache; -use crate::model::unet_model::UNetModel; +use crate::ml::model_cache::ModelCache; +use crate::ml::models::unet_model::UNetModel; // TODO: Clean up + static UNET_MODEL: Lazy>>> = Lazy::new(|| Arc::new(RwLock::new(None))); -pub struct Args { +// 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 api: Api, pub prompt: String, pub uncond_prompt: String, pub cpu: bool, @@ -30,9 +36,10 @@ pub struct Args { pub seed: Option, pub sd_version: StableDiffusionVersion, pub guidance_scale: Option, + pub model_cache: &'a ModelCache, } -pub fn run(args: Args) -> Result<()> { +pub fn run(args: Args<'_>) -> Result<()> { println!("Starting image generation with the following configuration:"); println!(" Model: {:?}", args.sd_version); println!(" Prompt: {}", args.prompt); @@ -72,35 +79,30 @@ 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); - 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 VAE ..."); + let vae_file = args.api.model(repo.to_string()).get("vae/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")?; + 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)?; @@ -153,30 +155,30 @@ pub fn run(args: Args) -> Result<()> { let latents = (latents * scheduler.init_noise_sigma())?; println!("Initial noise shape: {:?}", latents.shape()); - match UNET_MODEL.write() { - Err(err) => return Err(anyhow::Error::msg(err.to_string())), - Ok(mut maybe_model) => { - match &*maybe_model { - None => { - 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)?; - *maybe_model = Some(unet); - } - Some(_model) => { - println!("Model already exists"); - } - } - - } - } - + //match UNET_MODEL.write() { + // Err(err) => return Err(anyhow::Error::msg(err.to_string())), + // Ok(mut maybe_model) => { + // match &*maybe_model { + // None => { + // 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)?; + // *maybe_model = Some(unet); + // } + // Some(_model) => { + // println!("Model already exists"); + // } + // } + // + // } + //} + println!("Starting diffusion process..."); let mut latents = latents; - - - + + + for (timestep_index, ×tep) in timesteps.iter().enumerate() { println!("Processing step {}/{}", timestep_index + 1, timesteps.len()); let latent_model_input = latents.clone(); @@ -187,24 +189,12 @@ 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_MODEL.read() { - Err(err) => return Err(anyhow::Error::msg(err.to_string())), - Ok(maybe_model) => { - match &*maybe_model { - None => { - return Err(anyhow::Error::msg("Model does not exist")); - } - Some(unet) => { - noise_pred = unet.inference(&latent_model_input, timestep as f64, &text_embeddings)?; - } - } - } - } - - + + noise_pred = args.model_cache.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"); @@ -215,7 +205,7 @@ pub fn run(args: Args) -> Result<()> { // return Err(anyhow::anyhow!("UNet inference failed: {}", e)); // } //}; - + println!("Noise prediction shape: {:?}", noise_pred.shape()); println!("Applying scheduler step...");