diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 713ccd8..f4a018a 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -48,7 +48,7 @@ jobs: strategy: fail-fast: false matrix: - platform: [ubuntu-latest, macos-latest, windows-2022] + platform: [ubuntu-latest, macos-latest] #windows-2022] runs-on: ${{ matrix.platform }} steps: - uses: actions/checkout@v4 @@ -106,7 +106,7 @@ jobs: strategy: fail-fast: false matrix: - platform: [ubuntu-22.04, windows-2022] + platform: [ubuntu-22.04] #, windows-2022] runs-on: ${{ matrix.platform }} steps: - uses: actions/checkout@v4 diff --git a/cli/src/main.rs b/cli/src/main.rs index cc6052a..5f2035d 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -8,6 +8,7 @@ use diffusion_rs::{ DbCacheParamsBuilder, EasyCacheParamsBuilder, HiresParamsBuilder, PreviewType, SpectrumCacheParamsBuilder, UCacheParamsBuilder, Upscaler, gen_img, }, + modifier::lazily_load_params_from_disk, preset::{ Anima2Weight, AnimaWeight, ChromaRadianceWeight, ChromaWeight, DiffInstructStarWeight, ErnieImageWeight, Flux1MiniWeight, Flux1Weight, Flux2Klein4BWeight, Flux2Klein9BWeight, @@ -182,11 +183,12 @@ fn main() { } if args.low_vram { - model_config - .clip_on_cpu(true) - .vae_tiling(true) - .flash_attention(true) - .offload_params_to_cpu(true); + model_config.vae_tiling(true).flash_attention(true); + + let (new_config, new_model_config) = + lazily_load_params_from_disk((config, model_config))?; + config = new_config; + model_config = new_model_config; } match args.preview { diff --git a/src/api.rs b/src/api.rs index 18deeec..1cc15a2 100644 --- a/src/api.rs +++ b/src/api.rs @@ -101,7 +101,7 @@ pub enum DiffusionError { } #[non_exhaustive] -#[derive(Clone, Debug, strum::Display)] +#[derive(Clone, Debug, strum::Display, strum::EnumIter)] #[strum(serialize_all = "lowercase")] /// Backend devices pub enum BackendDevice { @@ -111,11 +111,12 @@ pub enum BackendDevice { METAL, GPU, AUTO, + DISK, DEFAULT, } #[non_exhaustive] -#[derive(Clone, Debug, strum::Display)] +#[derive(Clone, Debug, strum::Display, strum::EnumIter, PartialEq, Eq, Hash)] #[strum(serialize_all = "lowercase")] /// Module that can be bound to a specific [BackendDevice] pub enum Module { @@ -333,10 +334,6 @@ pub struct ModelConfig { #[builder(default = "false")] enable_mmap: bool, - /// Place the weights in RAM to save VRAM, and automatically load them into VRAM when needed - #[builder(default = "false")] - offload_params_to_cpu: bool, - /// Maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables graph splitting; -1 auto-detects free VRAM minus 1 GiB #[builder(default = "-1.0")] max_vram: f32, @@ -361,6 +358,10 @@ pub struct ModelConfig { #[builder(default = "Default::default()")] diffusion_model: CLibPath, + /// Path to the standalone unconditional diffusion model, currently used by Ideogram4 CFG + #[builder(default = "Default::default()")] + unconditional_diffusion_model: CLibPath, + /// Path to the qwen2vl text encoder #[builder(default = "Default::default()")] llm: CLibPath, @@ -463,18 +464,6 @@ pub struct ModelConfig { #[builder(default = "Prediction::PREDICTION_COUNT")] prediction: Prediction, - /// Keep vae in cpu (for low vram) (default: false) - #[builder(default = "false")] - vae_on_cpu: bool, - - /// keep clip in cpu (for low vram) (default: false) - #[builder(default = "false")] - clip_on_cpu: bool, - - /// Keep controlnet in cpu (for low vram) (default: false) - #[builder(default = "false")] - control_net_cpu: bool, - /// Use flash attention to reduce memory usage (model only). /// For most backends, it slows things down, but for cuda it generally speeds it up too. At the moment, it is only supported for some models and some backends (like cpu, cuda/rocm, metal). #[builder(default = "false")] @@ -756,7 +745,6 @@ impl ModelConfig { if self.upscaler_ctx.is_none() { let upscaler = new_upscaler_ctx( self.upscale_model.as_ref().unwrap().as_ptr(), - self.offload_params_to_cpu, self.diffusion_conv_direct, self.n_threads, self.upscale_tile_size, @@ -770,14 +758,12 @@ impl ModelConfig { } } - unsafe fn diffusion_ctx(&mut self, vae_decode_only: bool) -> *mut sd_ctx_t { + unsafe fn diffusion_ctx(&mut self) -> *mut sd_ctx_t { unsafe { // This is required to support img2img after text2img generation // otherwise the context is cached and won't have a decode graph // leading to an assertion error in sdcpp - if let Some((sd_ctx, sd_ctx_params)) = self.diffusion_ctx.as_ref() - && sd_ctx_params.vae_decode_only != vae_decode_only - { + if let Some((sd_ctx, _)) = self.diffusion_ctx.as_ref() { sd_set_progress_callback(None, null_mut()); free_sd_ctx(*sd_ctx); self.diffusion_ctx = None; @@ -793,20 +779,16 @@ impl ModelConfig { high_noise_diffusion_model_path: self.high_noise_diffusion_model.as_ptr(), t5xxl_path: self.t5xxl.as_ptr(), diffusion_model_path: self.diffusion_model.as_ptr(), + uncond_diffusion_model_path: self.unconditional_diffusion_model.as_ptr(), vae_path: self.vae.as_ptr(), taesd_path: self.taesd.as_ptr(), control_net_path: self.control_net.as_ptr(), embeddings: self.embeddings.2.as_ptr(), embedding_count: self.embeddings.1.len() as u32, photo_maker_path: self.photo_maker.as_ptr(), - vae_decode_only, - free_params_immediately: false, n_threads: self.n_threads, wtype: self.weight_type, rng_type: self.rng, - keep_clip_on_cpu: self.clip_on_cpu, - keep_control_net_on_cpu: self.control_net_cpu, - keep_vae_on_cpu: self.vae_on_cpu, diffusion_flash_attn: self.diffusion_flash_attention, flash_attn: self.flash_attention, diffusion_conv_direct: self.diffusion_conv_direct, @@ -814,7 +796,6 @@ impl ModelConfig { chroma_use_t5_mask: self.chroma_enable_t5_mask, chroma_t5_mask_pad: self.chroma_t5_mask_pad, vae_conv_direct: self.vae_conv_direct, - offload_params_to_cpu: self.offload_params_to_cpu, prediction: self.prediction, force_sdxl_vae_conv_scale: self.force_sdxl_vae_conv_scale, tae_preview_only: self.taesd_preview_only, @@ -867,10 +848,10 @@ impl From<&ModelConfig> for ModelConfigBuilder { builder .n_threads(value.n_threads) .max_vram(value.max_vram) - .offload_params_to_cpu(value.offload_params_to_cpu) .upscale_repeats(value.upscale_repeats) .model(value.model.clone()) .diffusion_model(value.diffusion_model.clone()) + .unconditional_diffusion_model(value.unconditional_diffusion_model.clone()) .llm(value.llm.clone()) .llm_vision(value.llm_vision.clone()) .clip_l(value.clip_l.clone()) @@ -894,10 +875,7 @@ impl From<&ModelConfig> for ModelConfigBuilder { .scheduler(value.scheduler) .sigmas(value.sigmas.clone()) .prediction(value.prediction) - .vae_on_cpu(value.vae_on_cpu) - .clip_on_cpu(value.clip_on_cpu) .control_net(value.control_net.clone()) - .control_net_cpu(value.control_net_cpu) .flash_attention(value.flash_attention) .chroma_disable_dit_mask(value.chroma_disable_dit_mask) .chroma_enable_t5_mask(value.chroma_enable_t5_mask) @@ -1340,8 +1318,7 @@ fn gen_img_maybe_progress( let has_init_image = config.init_img.exists(); let has_mask_image = config.mask_img.exists(); - let is_decode_only = !has_init_image; - let sd_ctx = model_config.diffusion_ctx(is_decode_only); + let sd_ctx = model_config.diffusion_ctx(); let upscaler_ctx = model_config.upscaler_ctx(); let mut init_image = sd_image_t { diff --git a/src/modifier.rs b/src/modifier.rs index 67b1529..98b940c 100644 --- a/src/modifier.rs +++ b/src/modifier.rs @@ -1,7 +1,10 @@ +use std::collections::HashMap; + use hf_hub::api::sync::ApiError; +use strum::IntoEnumIterator; use crate::{ - api::{LoraSpec, PreviewType, SampleMethod}, + api::{BackendDevice, LoraSpec, Module, PreviewType, SampleMethod}, preset::ConfigsBuilder, util::download_file_hf_hub, }; @@ -266,7 +269,23 @@ pub fn t5xxl_q8_0_flux_1(mut builder: ConfigsBuilder) -> Result Result { - builder.1.offload_params_to_cpu(true); + let params: HashMap<_, _> = Module::iter() + .map(|module| (module, BackendDevice::CPU)) + .collect(); + // params. + builder.1.params_backend(params); + Ok(builder) +} + +/// Lazily load model parameters from disk (for low VRAM GPUs) +pub fn lazily_load_params_from_disk( + mut builder: ConfigsBuilder, +) -> Result { + let params: HashMap<_, _> = Module::iter() + .map(|module| (module, BackendDevice::DISK)) + .collect(); + // params. + builder.1.params_backend(params); Ok(builder) } diff --git a/src/preset_builder.rs b/src/preset_builder.rs index 989eea5..e265cb8 100644 --- a/src/preset_builder.rs +++ b/src/preset_builder.rs @@ -3,8 +3,8 @@ use std::path::PathBuf; use crate::{ api::{ModelConfigBuilder, SampleMethod, Scheduler}, modifier::{ - sdxl_vae_fp16_fix, t5xxl_fp16_flux_1, t5xxl_q2_k_flux_1, t5xxl_q3_k_flux_1, - t5xxl_q4_k_flux_1, t5xxl_q8_0_flux_1, + offload_params_to_cpu, sdxl_vae_fp16_fix, t5xxl_fp16_flux_1, t5xxl_q2_k_flux_1, + t5xxl_q3_k_flux_1, t5xxl_q4_k_flux_1, t5xxl_q8_0_flux_1, }, preset::{ Anima2Weight, AnimaWeight, ChromaRadianceWeight, ChromaWeight, ConfigsBuilder, @@ -494,7 +494,6 @@ pub fn flux_2_dev(sd_type: Flux2Weight) -> Result { model_config .diffusion_model(model) .llm(llm) - .offload_params_to_cpu(true) .flash_attention(true) .vae(vae) .vae_tiling(true); @@ -502,7 +501,7 @@ pub fn flux_2_dev(sd_type: Flux2Weight) -> Result { .cfg_scale(1.) .sampling_method(SampleMethod::EULER_SAMPLE_METHOD); - Ok((config, model_config)) + offload_params_to_cpu((config, model_config)) } fn flux_2_dev_weight(sd_type: Flux2Weight) -> Result<(PathBuf, PathBuf), ApiError> { @@ -690,7 +689,6 @@ pub fn qwen_image(sd_type: QwenImageWeight) -> Result .diffusion_model(model) .llm(llm) .vae(vae) - .offload_params_to_cpu(true) .flash_attention(true) .vae_tiling(true) .flow_shift(3.0); @@ -700,7 +698,7 @@ pub fn qwen_image(sd_type: QwenImageWeight) -> Result .height(1024) .width(1024); - Ok((config, model_config)) + offload_params_to_cpu((config, model_config)) } fn qwen_image_weight(sd_type: QwenImageWeight) -> Result<(PathBuf, PathBuf), ApiError> { @@ -814,12 +812,11 @@ pub fn ovis_image(sd_type: OvisImageWeight) -> Result .diffusion_model(model) .llm(llm) .vae(vae) - .offload_params_to_cpu(true) .flash_attention(true) .vae_tiling(true); config.steps(20).cfg_scale(5.).height(512).width(512); - Ok((config, model_config)) + offload_params_to_cpu((config, model_config)) } fn ovis_image_weight(sd_type: OvisImageWeight) -> Result<(PathBuf, PathBuf), ApiError> { @@ -990,12 +987,11 @@ pub fn flux_2_klein_4b(sd_type: Flux2Klein4BWeight) -> Result Result<(PathBuf, PathBuf), ApiError> { @@ -1034,12 +1030,11 @@ pub fn flux_2_klein_base_4b(sd_type: Flux2KleinBase4BWeight) -> Result Result Result<(PathBuf, PathBuf), ApiError> { @@ -1130,12 +1124,11 @@ pub fn flux_2_klein_base_9b(sd_type: Flux2KleinBase9BWeight) -> Result