From aa25ae8c87f5d8d43daca6fa803c1036008b8b2a Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 27 Mar 2025 21:04:39 -0400 Subject: [PATCH] load flux models from R2 --- .../ml_models/src/ml/flux/flux_pipeline.rs | 162 +++++++++--------- crates/desktop/ml_models/src/ml/flux/mod.rs | 4 +- .../src/endpoints/text_to_image_endpoint.rs | 19 +- 3 files changed, 97 insertions(+), 88 deletions(-) diff --git a/crates/desktop/ml_models/src/ml/flux/flux_pipeline.rs b/crates/desktop/ml_models/src/ml/flux/flux_pipeline.rs index 9afbfa72b7..2344cbf25c 100644 --- a/crates/desktop/ml_models/src/ml/flux/flux_pipeline.rs +++ b/crates/desktop/ml_models/src/ml/flux/flux_pipeline.rs @@ -1,3 +1,4 @@ +use std::path::Path; use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}; use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor; @@ -11,7 +12,7 @@ use candle_core::{DType, IndexOp, Tensor, D, Device, Module}; use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution}; use candle_transformers::models::stable_diffusion::lcm::LCMScheduler; use image::DynamicImage; -use log::info; +use log::{debug, info}; use rand::Rng; use tauri::{AppHandle, Emitter}; use candle_transformers::models::{clip, t5, flux}; @@ -38,10 +39,41 @@ pub struct FluxPipelineArgs<'a> { pub seed: Option, pub use_quantized_model: Option, pub model_config: &'a ModelConfig, + pub flux_dev_or_schnell_path: &'a Path, + pub flux_schnell_quantized_gguf: &'a Path, + pub flux_autoencoder_path: &'a Path, + pub google_t5_model_path: &'a Path, + pub google_t5_config_path: &'a Path, + pub lmz_t5_tokenizer_path: &'a Path, + pub openai_clip_model_path: &'a Path, + pub openai_clip_tokenizer_json_path: &'a Path, } pub fn flux_pipeline(args: FluxPipelineArgs<'_>) -> Result { - let FluxPipelineArgs { prompt, model_cache, prompt_cache, app, image, flux_model, active_blocks, use_cuda_stream, prefetch_next_batch, use_full_gpu, img2img_strength, seed, use_quantized_model, model_config } = args; + let FluxPipelineArgs { + prompt, + model_cache, + prompt_cache, + app, + image, + flux_model, + active_blocks, + use_cuda_stream, + prefetch_next_batch, + use_full_gpu, + img2img_strength, + seed, + use_quantized_model, + model_config , + flux_dev_or_schnell_path, + flux_schnell_quantized_gguf, + flux_autoencoder_path, + google_t5_model_path, + google_t5_config_path, + lmz_t5_tokenizer_path, + openai_clip_model_path, + openai_clip_tokenizer_json_path, + } = args; if use_quantized_model.unwrap_or(false) { #[cfg(feature = "cuda")] @@ -55,14 +87,15 @@ pub fn flux_pipeline(args: FluxPipelineArgs<'_>) -> Result { let device = model_config.device.clone(); let dtype = device.bf16_default_to_f32(); - println!("Starting image generation with FLUX model:"); - println!(" Model: {:?}", flux_model); - println!(" Prompt: {}", prompt); - println!(" Dimensions: {}x{}", width, height); - println!(" Device: {:?}", model_config.device); + info!("Starting image generation with FLUX model:"); + info!(" Model: {:?}", flux_model); + info!(" Prompt: {}", prompt); + info!(" Dimensions: {}x{}", width, height); + info!(" Device: {:?}", model_config.device); let seed_value = seed.unwrap_or_else(|| rand::thread_rng().gen()); device.set_seed(seed_value)?; + info!("Using seed: {}", seed_value); // Setup FLUX model configuration @@ -71,48 +104,30 @@ pub fn flux_pipeline(args: FluxPipelineArgs<'_>) -> Result { FluxModel::Schnell => flux::model::Config::schnell(), }; - let bf_repo = match flux_model { - FluxModel::Dev => model_config.hf_api.repo(hf_hub::Repo::model("black-forest-labs/FLUX.1-dev".to_string())), - FluxModel::Schnell => model_config.hf_api.repo(hf_hub::Repo::model("black-forest-labs/FLUX.1-schnell".to_string())), - }; // Get T5 embeddings - println!("Generating T5 embeddings for prompt..."); let t5_emb = { - let repo = model_config.hf_api.repo(hf_hub::Repo::with_revision("google/t5-v1_1-xxl".to_string(), hf_hub::RepoType::Model, "refs/pr/2".to_string())); - println!("Downloading model (1)"); - let model_file = repo.get("model.safetensors")?; - println!("model downloaded (1)"); - let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype , &device)? }; - println!("Downloading model (2)"); - let config_filename = repo.get("config.json")?; - println!("model downloaded (2)"); - let config = std::fs::read_to_string(config_filename)?; + info!("Generating T5 embeddings for prompt..."); + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[google_t5_model_path], dtype , &device)? }; + let config = std::fs::read_to_string(google_t5_config_path)?; let config: t5::Config = serde_json::from_str(&config)?; - let mut model = t5::T5EncoderModel::load(vb, &config)?; - println!("Downloading model (3)"); - let tokenizer_filename = model_config.hf_api.model("lmz/mt5-tokenizers".to_string()).get("t5-v1_1-xxl.tokenizer.json")?; - println!("model downloaded (3)"); - let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?; - println!("encoding"); + let mut model = t5::T5EncoderModel::load(vb, &config)?; + let tokenizer = Tokenizer::from_file(lmz_t5_tokenizer_path).map_err(E::msg)?; + let mut tokens = tokenizer.encode(prompt.as_str(), true).map_err(E::msg)?.get_ids().to_vec(); tokens.resize(256, 0); let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?; - println!("forward"); model.forward(&input_token_ids)? }; - println!("T5 embeddings generated"); + + info!("T5 embeddings generated"); // Get CLIP embeddings - println!("Generating CLIP embeddings for prompt..."); let clip_emb = { - let repo = model_config.hf_api.repo(hf_hub::Repo::model( - "openai/clip-vit-large-patch14".to_string(), - )); - let model_file = repo.get("model.safetensors")?; + info!("Generating CLIP embeddings for prompt..."); let vb = - unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? }; + unsafe { VarBuilder::from_mmaped_safetensors(&[openai_clip_model_path], dtype, &device)? }; // https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json let config = clip::text_model::ClipTextConfig { vocab_size: 49408, @@ -127,37 +142,34 @@ pub fn flux_pipeline(args: FluxPipelineArgs<'_>) -> Result { }; let model = clip::text_model::ClipTextTransformer::new(vb.pp("text_model"), &config)?; - let tokenizer_filename = repo.get("tokenizer.json")?; - let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?; + let tokenizer = Tokenizer::from_file(openai_clip_tokenizer_json_path).map_err(E::msg)?; let tokens = tokenizer .encode(prompt.as_str(), true) .map_err(E::msg)? .get_ids() .to_vec(); let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?; - - model.forward(&input_token_ids)? }; - println!("CLIP embeddings generated"); + info!("CLIP embeddings generated"); // Process input image to tensor - println!("Converting input image to tensor..."); + info!("Converting input image to tensor..."); let img_tensor = dynamic_image_to_tensor(image, &device, dtype)?; - println!("Image tensor shape: {:?}", img_tensor.shape()); + info!("Image tensor shape: {:?}", img_tensor.shape()); // Create initial noise - this is important to match the expected format - println!("Generating initial noise for FLUX..."); + info!("Generating initial noise for FLUX..."); let initial_tensor = flux::sampling::get_noise(1, height, width, &device)?.to_dtype(dtype)?; - println!("Initial tensor shape: {:?}", initial_tensor.shape()); + info!("Initial tensor shape: {:?}", initial_tensor.shape()); // Blend initial noise with input image if img2img_strength is specified let initial_tensor = if let Some(strength) = img2img_strength { if strength > 0.0 && strength < 1.0 { - println!("Using img2img with strength: {}", strength); + info!("Using img2img with strength: {}", strength); // Reshape the input tensor to match the noise tensor dimensions let adjusted_img = if img_tensor.shape() != initial_tensor.shape() { - println!("Reshaping input image to match noise dimensions..."); + info!("Reshaping input image to match noise dimensions..."); // You may need more complex reshaping logic here depending on your use case let (_, c, h, w) = initial_tensor.dims4()?; img_tensor.upsample_nearest2d(h, w)? @@ -177,10 +189,10 @@ pub fn flux_pipeline(args: FluxPipelineArgs<'_>) -> Result { } else { initial_tensor }; - println!("Final initial tensor shape: {:?}", initial_tensor.shape()); + info!("Final initial tensor shape: {:?}", initial_tensor.shape()); // Create FLUX state - println!("Creating FLUX state..."); + info!("Creating FLUX state..."); let state = if quantized { flux::sampling::State::new( &t5_emb.to_dtype(DType::F32)?, @@ -190,7 +202,7 @@ pub fn flux_pipeline(args: FluxPipelineArgs<'_>) -> Result { } else { flux::sampling::State::new(&t5_emb, &clip_emb, &initial_tensor)? }; - println!("{state:?}"); + info!("{state:?}"); // Set appropriate timesteps based on model let timesteps = match flux_model { @@ -198,98 +210,88 @@ pub fn flux_pipeline(args: FluxPipelineArgs<'_>) -> Result { FluxModel::Schnell => flux::sampling::get_schedule(4, None), }; - println!("Starting FLUX denoising with timesteps: {:?}", timesteps); + info!("Starting FLUX denoising with timesteps: {:?}", timesteps); // Actual model loading and inference let img_latent = if quantized { // Quantized model path - println!("Loading quantized FLUX model..."); + info!("Loading quantized FLUX model..."); let model_file = match flux_model { - FluxModel::Schnell => model_config.hf_api.repo(hf_hub::Repo::model("lmz/candle-flux".to_string())).get("flux1-schnell.gguf")?, + FluxModel::Schnell => flux_schnell_quantized_gguf, FluxModel::Dev => return Err(anyhow!("Quantized model not available for FLUX.1-dev")), }; - println!("Loading quantized model from {}", model_file.to_str().unwrap()); + info!("Loading quantized model from {}", model_file.to_str().unwrap()); let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(model_file, &device)?; let mut model = flux::quantized_model::Flux::new(&flux_config, vb)?; - println!("Quantized model loaded on {:?}", device); - + info!("Quantized model loaded on {:?}", device); // Run denoising - println!("Running denoising..."); + info!("Running denoising..."); let img_latent = flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, ×teps, 4.)?.to_dtype(dtype)?; - println!("Denoising completed"); + info!("Denoising completed"); img_latent } else { // Full precision model - println!("Loading FLUX model..."); - let model_file = match flux_model { - FluxModel::Schnell => bf_repo.get("flux1-schnell.safetensors")?, - FluxModel::Dev => bf_repo.get("flux1-dev.safetensors")?, - }; - + info!("Loading FLUX model..."); let mut model = if use_full_gpu.unwrap_or(false) { // Load entire model on GPU - println!("Loading entire model directly to GPU..."); - let gpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? }; + info!("Loading entire model directly to GPU..."); + let gpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[flux_dev_or_schnell_path], dtype, &device)? }; let mut model = flux::model::Flux::new(&flux_config, gpu_vb)?; model.set_use_device_management(false); model } else { // Use memory-efficient approach with blocks on CPU - println!("Using memory-efficient approach with blocks on CPU..."); + info!("Using memory-efficient approach with blocks on CPU..."); let cpu_device = candle_core::Device::Cpu; - let cpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &cpu_device)? }; + let cpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[flux_dev_or_schnell_path], dtype, &cpu_device)? }; let mut model = flux::model::Flux::new_with_gpu_core(&flux_config, cpu_vb, &device)?; // Set active blocks let active_blocks_count = active_blocks.unwrap_or(1); model.set_max_active_blocks(active_blocks_count); - println!("Using {} active blocks on GPU", active_blocks_count); + info!("Using {} active blocks on GPU", active_blocks_count); // Enable prefetching if requested if prefetch_next_batch.unwrap_or(false) { model.set_prefetch_next_batch(true); - println!("Prefetching enabled - will overlap transfers with computation"); + info!("Prefetching enabled - will overlap transfers with computation"); } model }; - // Run denoising flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, ×teps, 4.)? }; - println!("FLUX denoising completed, unpacking latent image..."); + info!("FLUX denoising completed, unpacking latent image..."); let img_unpacked = flux::sampling::unpack(&img_latent, height, width)?; // Load autoencoder for final decoding - println!("Loading autoencoder for final decoding..."); + info!("Loading autoencoder for final decoding..."); + let ae_config = match flux_model { FluxModel::Dev => flux::autoencoder::Config::dev(), FluxModel::Schnell => flux::autoencoder::Config::schnell(), }; - let model_file = bf_repo.get("ae.safetensors")?; - - - let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? }; + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[flux_autoencoder_path], dtype, &device)? }; let ae_model = flux::autoencoder::AutoEncoder::new(&ae_config, vb)?; - - println!("Decoding final image..."); + info!("Decoding final image..."); let img = ae_model.decode(&img_unpacked)?; // Convert tensor to image - println!("Converting tensor to image..."); + info!("Converting tensor to image..."); let img = ((img.clamp(-1f32, 1f32)? + 1.0)? * 127.5)?.to_dtype(candle_core::DType::U8)?; // Convert to RGB image let image = tensor_to_image_buffer(&img.i(0)?)?; - println!("FLUX pipeline completed successfully"); + info!("FLUX pipeline completed successfully"); Ok(image) } diff --git a/crates/desktop/ml_models/src/ml/flux/mod.rs b/crates/desktop/ml_models/src/ml/flux/mod.rs index ededd1ec2f..b960caac14 100644 --- a/crates/desktop/ml_models/src/ml/flux/mod.rs +++ b/crates/desktop/ml_models/src/ml/flux/mod.rs @@ -1,3 +1 @@ -pub mod flux_pipeline; - -pub use flux_pipeline::{flux_pipeline, FluxPipelineArgs, FluxModel}; \ No newline at end of file +pub mod flux_pipeline; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs index e16cb5db5d..53ca93157d 100644 --- a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs @@ -1,22 +1,20 @@ #[cfg(not(target_os = "macos"))] -use ml_models::ml::flux::flux_pipeline::{flux_pipeline, FluxPipelineArgs, FluxModel}; +use ml_models::ml::flux::flux_pipeline::{flux_pipeline, FluxModel, FluxPipelineArgs}; use crate::state::app_config::AppConfig; -use crate::state::app_dir::AppDataRoot; +use crate::state::app_dir::{AppDataRoot, AppWeightsDir}; use crate::stubs::model_cache::ModelCache; use crate::stubs::prompt_cache::PromptCache; use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png; use image::imageops::FilterType; use image::{DynamicImage, ImageReader, RgbImage}; use log::{error, info}; -use ml_weights_registry::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; +use ml_weights_registry::weights_registry::flux_weights::{FLUX_DEV, FLUX_SCHNELL, FLUX_SCHNELL_AUTOENCODER, GOOGLE_T5_V1_1_XXL_CONFIG, GOOGLE_T5_V1_1_XXL_MODEL, LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF, LMZ_T5_TOKENIZER_JSON, OPENAI_CLIP_VIT_P14_MODEL, OPENAI_CLIP_VIT_P14_TOKENIZER_JSON}; use once_cell::sync::Lazy; use std::io::Cursor; use tauri::{AppHandle, State}; -const RANDOM_SEED: u32 = 42; - const STRENGTH: f64 = 100.0; const PNG_BYTES : &[u8] = include_bytes!("../../binary_includes/1024.png"); @@ -36,6 +34,7 @@ pub async fn text_to_image( model_cache: State<'_, ModelCache>, prompt_cache: State<'_, PromptCache>, app: AppHandle, + app_data_root: State<'_, AppDataRoot>, ) -> Result { info!("text_to_image endpoint called."); @@ -50,6 +49,7 @@ pub async fn text_to_image( &model_config, &model_cache, prompt_cache, + app_data_root.weights_dir(), app, ).await; @@ -78,6 +78,7 @@ async fn text_to_image_impl( config: &AppConfig, model_cache: &ModelCache, prompt_cache: State<'_, PromptCache>, + weights_dir: &AppWeightsDir, app: AppHandle, ) -> Result { @@ -96,6 +97,14 @@ async fn text_to_image_impl( seed: None, use_quantized_model: Some(true), model_config: &config.model_config, + flux_dev_or_schnell_path: &weights_dir.weight_path(&FLUX_SCHNELL), + flux_schnell_quantized_gguf: &weights_dir.weight_path(&LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF), + flux_autoencoder_path: &weights_dir.weight_path(&FLUX_SCHNELL_AUTOENCODER), + google_t5_model_path: &weights_dir.weight_path(&GOOGLE_T5_V1_1_XXL_MODEL), + google_t5_config_path: &weights_dir.weight_path(&GOOGLE_T5_V1_1_XXL_CONFIG), + lmz_t5_tokenizer_path: &weights_dir.weight_path(&LMZ_T5_TOKENIZER_JSON), + openai_clip_model_path: &weights_dir.weight_path(&OPENAI_CLIP_VIT_P14_MODEL), + openai_clip_tokenizer_json_path: &weights_dir.weight_path(&OPENAI_CLIP_VIT_P14_TOKENIZER_JSON), }; let image = flux_pipeline(args)