feat: implement lazy loading of model parameters from disk

.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

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 {

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,28 +779,23 @@ 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,
                     chroma_use_dit_mask: !self.chroma_disable_dit_mask,
                     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 {

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<ConfigsBuilder,
 
 /// Offload model parameters to CPU (for low VRAM GPUs)
 pub fn offload_params_to_cpu(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
-    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<ConfigsBuilder, ApiError> {
+    let params: HashMap<_, _> = Module::iter()
+        .map(|module| (module, BackendDevice::DISK))
+        .collect();
+    // params.
+    builder.1.params_backend(params);
     Ok(builder)
 }
 

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,15 +494,14 @@ pub fn flux_2_dev(sd_type: Flux2Weight) -> Result<ConfigsBuilder, ApiError> {
     model_config
         .diffusion_model(model)
         .llm(llm)
-        .offload_params_to_cpu(true)
         .flash_attention(true)
         .vae(vae)
         .vae_tiling(true);
     config
         .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<ConfigsBuilder, ApiError>
         .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<ConfigsBuilder, ApiError>
         .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<ConfigsBuilder, ApiError>
         .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<ConfigsBuilder, Ap
         .diffusion_model(model)
         .llm(llm)
         .vae(vae)
-        .offload_params_to_cpu(true)
         .flash_attention(true)
         .vae_tiling(true);
     config.cfg_scale(1.).steps(4).height(1024).width(1024);
 
-    Ok((config, model_config))
+    offload_params_to_cpu((config, model_config))
 }
 
 fn flux_2_klein_4b_weight(sd_type: Flux2Klein4BWeight) -> Result<(PathBuf, PathBuf), ApiError> {
@@ -1034,12 +1030,11 @@ pub fn flux_2_klein_base_4b(sd_type: Flux2KleinBase4BWeight) -> Result<ConfigsBu
         .diffusion_model(model)
         .llm(llm)
         .vae(vae)
-        .offload_params_to_cpu(true)
         .flash_attention(true)
         .vae_tiling(true);
     config.cfg_scale(4.).steps(20).height(1024).width(1024);
 
-    Ok((config, model_config))
+    offload_params_to_cpu((config, model_config))
 }
 
 fn flux_2_klein_base_4b_weight(
@@ -1086,12 +1081,11 @@ pub fn flux_2_klein_9b(sd_type: Flux2Klein9BWeight) -> Result<ConfigsBuilder, Ap
         .diffusion_model(model)
         .llm(llm)
         .vae(vae)
-        .offload_params_to_cpu(true)
         .flash_attention(true)
         .vae_tiling(true);
     config.cfg_scale(1.).steps(4).height(1024).width(1024);
 
-    Ok((config, model_config))
+    offload_params_to_cpu((config, model_config))
 }
 
 fn flux_2_klein_9b_weight(sd_type: Flux2Klein9BWeight) -> Result<(PathBuf, PathBuf), ApiError> {
@@ -1130,12 +1124,11 @@ pub fn flux_2_klein_base_9b(sd_type: Flux2KleinBase9BWeight) -> Result<ConfigsBu
         .diffusion_model(model)
         .llm(llm)
         .vae(vae)
-        .offload_params_to_cpu(true)
         .flash_attention(true)
         .vae_tiling(true);
     config.cfg_scale(4.).steps(20).height(1024).width(1024);
 
-    Ok((config, model_config))
+    offload_params_to_cpu((config, model_config))
 }
 
 fn flux_2_klein_base_9b_weight(

sys/Cargo.toml

@@ -16,15 +16,45 @@ include = [
     "stable-diffusion.cpp/src/*.h",
     "stable-diffusion.cpp/src/*.hpp",
     "stable-diffusion.cpp/src/*.cpp",
+    "stable-diffusion.cpp/src/conditioning/*.h",
+    "stable-diffusion.cpp/src/conditioning/*.hpp",
+    "stable-diffusion.cpp/src/conditioning/*.cpp",
+    "stable-diffusion.cpp/src/core/*.h",
+    "stable-diffusion.cpp/src/core/*.hpp",
+    "stable-diffusion.cpp/src/core/*.cpp",
+    "stable-diffusion.cpp/src/extensions/*.h",
+    "stable-diffusion.cpp/src/extensions/*.hpp",
+    "stable-diffusion.cpp/src/extensions/*.cpp",
+    "stable-diffusion.cpp/src/model/adapter/*.h",
+    "stable-diffusion.cpp/src/model/adapter/*.hpp",
+    "stable-diffusion.cpp/src/model/adapter/*.cpp",
+    "stable-diffusion.cpp/src/model/common/*.h",
+    "stable-diffusion.cpp/src/model/common/*.hpp",
+    "stable-diffusion.cpp/src/model/common/*.cpp",
+    "stable-diffusion.cpp/src/model/diffusion/*.h",
+    "stable-diffusion.cpp/src/model/diffusion/*.hpp",
+    "stable-diffusion.cpp/src/model/diffusion/*.cpp",
+    "stable-diffusion.cpp/src/model/te/*.h",
+    "stable-diffusion.cpp/src/model/te/*.hpp",
+    "stable-diffusion.cpp/src/model/te/*.cpp",
+    "stable-diffusion.cpp/src/model/upscaler/*.h",
+    "stable-diffusion.cpp/src/model/upscaler/*.hpp",
+    "stable-diffusion.cpp/src/model/upscaler/*.cpp",
+    "stable-diffusion.cpp/src/model/vae/*.h",
+    "stable-diffusion.cpp/src/model/vae/*.hpp",
+    "stable-diffusion.cpp/src/model/vae/*.cpp",
+    "stable-diffusion.cpp/src/model_io/*.h",
+    "stable-diffusion.cpp/src/model_io/*.hpp",
+    "stable-diffusion.cpp/src/model_io/*.cpp",
+    "stable-diffusion.cpp/src/runtime/*.h",
+    "stable-diffusion.cpp/src/runtime/*.hpp",
+    "stable-diffusion.cpp/src/runtime/*.cpp",
     "stable-diffusion.cpp/src/tokenizers/*.h",
     "stable-diffusion.cpp/src/tokenizers/*.hpp",
     "stable-diffusion.cpp/src/tokenizers/*.cpp",
     "stable-diffusion.cpp/src/tokenizers/vocab/*.h",
     "stable-diffusion.cpp/src/tokenizers/vocab/*.hpp",
     "stable-diffusion.cpp/src/tokenizers/vocab/*.cpp",
-    "stable-diffusion.cpp/src/model_io/*.h",
-    "stable-diffusion.cpp/src/model_io/*.hpp",
-    "stable-diffusion.cpp/src/model_io/*.cpp",
     "stable-diffusion.cpp/thirdparty/**",
     "stable-diffusion.cpp/ggml/cmake/**",
     "stable-diffusion.cpp/ggml/src/*",