feat: unet spectrum caching support

cli/README.md

@@ -6,10 +6,10 @@ cli.exe flux1schnell  "a lovely duck holding a sign says 'drink your water'" --r
 ```
 ## Help
 ```
-Usage: cli.exe [OPTIONS] --preset <PRESET> <PROMPT>
+Usage: diffusion-rs-cli.exe [OPTIONS] <PRESET> <PROMPT>
 
 Arguments:
-  <PRESET>  The preset to use [possible values: StableDiffusion1_4, StableDiffusion1_5, StableDiffusion2_1, StableDiffusion3Medium, StableDiffusion3_5Medium, StableDiffusion3_5Large, StableDiffusion3_5LargeTurbo, SDXLBase1_0, SDTurbo, SDXLTurbo1_0, Flux1Dev, Flux1Schnell, Flux1Mini, JuggernautXL11, Chroma, NitroSDRealism, NitroSDVibrant, DiffInstructStar, ChromaRadiance, SSD1B, Flux2Dev, ZImageTurbo, QwenImage, OvisImage, DreamShaperXL2_1Turbo, TwinFlowZImageTurboExp, SDXS512DreamShaper, Flux2Klein4B, Flux2KleinBase4B, Flux2Klein9B, Flux2KleinBase9B, SegmindVega]
+  <PRESET>  The preset to use [possible values: StableDiffusion1_4, StableDiffusion1_5, StableDiffusion2_1, StableDiffusion3Medium, StableDiffusion3_5Medium, StableDiffusion3_5Large, StableDiffusion3_5LargeTurbo, SDXLBase1_0, SDTurbo, SDXLTurbo1_0, Flux1Dev, Flux1Schnell, Flux1Mini, JuggernautXL11, Chroma, NitroSDRealism, NitroSDVibrant, DiffInstructStar, ChromaRadiance, SSD1B, Flux2Dev, ZImageTurbo, QwenImage, OvisImage, DreamShaperXL2_1Turbo, TwinFlowZImageTurboExp, SDXS512DreamShaper, Flux2Klein4B, Flux2KleinBase4B, Flux2Klein9B, Flux2KleinBase9B, SegmindVega, Anima]
   <PROMPT>  The prompt to render
 
 Options:
@@ -20,10 +20,11 @@ Options:
       --height <HEIGHT>      Override the preset default height
   -b, --batch <BATCH>        Number of images to generate (default 1) [default: 1]
   -o, --output <OUTPUT>      Output Folder [default: ./]
+  -c, --cache <CACHE>        Caching methods accelerate diffusion inference [possible values: ucache, easycache, dbcache, taylorseer, cachedit, spectrum]    
   -p, --preview <PREVIEW>    Enable preview [possible values: fast, accurate]
   -t, --token <TOKEN>        Set Huggingface Hub token. Only used when downloading models that have not been cached before
   -l, --low-vram             Enable optimization to use less VRAM: clip_on_cpu, vae tiling, flash_attention, offload_params_to_cpu
-  -r, --random-seed          Enable Random Seed: different runs will produce different results      
+  -r, --random-seed          Enable Random Seed: different runs will produce different results
   -h, --help                 Print help
   -V, --version              Print version
   ```

cli/src/main.rs

@@ -5,7 +5,8 @@ use clap::{Parser, ValueEnum};
 
 use diffusion_rs::{
     api::{
-        DbCacheParamsBuilder, EasyCacheParamsBuilder, PreviewType, UCacheParamsBuilder, gen_img,
+        DbCacheParamsBuilder, EasyCacheParamsBuilder, PreviewType, SpectrumCacheParamsBuilder,
+        UCacheParamsBuilder, gen_img,
     },
     preset::{
         AnimaWeight, ChromaRadianceWeight, ChromaWeight, DiffInstructStarWeight, Flux1MiniWeight,
@@ -33,6 +34,7 @@ enum CacheMode {
     DBCACHE,
     TAYLORSEER,
     CACHEDIT,
+    SPECTRUM,
 }
 
 #[derive(Clone, Debug, ValueEnum)]
@@ -80,7 +82,7 @@ struct Args {
     output: PathBuf,
 
     /// Caching methods accelerate diffusion inference
-    #[arg(long)]
+    #[arg(short, long)]
     cache: Option<CacheMode>,
 
     /// Enable preview
@@ -122,73 +124,76 @@ fn main() {
     }
     println!();
 
-    let (config, mut model_config) =
-        PresetBuilder::default()
-            .preset(preset)
-            .prompt(&args.prompt)
-            .with_modifier(move |(mut config, mut model_config)| {
-                // atch request?
-                if args.batch > 1 {
-                    config.batch_count(args.batch);
-                    config.output(args.output);
-                } else {
-                    config.output(file_name);
-                }
-
-                if let Some(steps) = &args.steps {
-                    config.steps(*steps);
-                }
-
-                if args.random_seed {
-                    config.seed(-1);
-                }
-
-                if let Some(width) = args.width {
-                    config.width(width);
-                }
-
-                if let Some(height) = args.height {
-                    config.height(height);
-                }
-
-                if let Some(negative) = args.negative {
-                    config.negative_prompt(negative);
-                }
-
-                if args.low_vram {
-                    model_config
-                        .clip_on_cpu(true)
-                        .vae_tiling(true)
-                        .flash_attention(true)
-                        .offload_params_to_cpu(true);
-                }
-
-                match args.preview {
-                    Some(PreviewMode::Fast) => config.preview_mode(PreviewType::PREVIEW_PROJ),
-                    Some(PreviewMode::Accurate) => config.preview_mode(PreviewType::PREVIEW_VAE),
-                    None => config.preview_mode(PreviewType::PREVIEW_NONE),
+    let (config, mut model_config) = PresetBuilder::default()
+        .preset(preset)
+        .prompt(&args.prompt)
+        .with_modifier(move |(mut config, mut model_config)| {
+            // atch request?
+            if args.batch > 1 {
+                config.batch_count(args.batch);
+                config.output(args.output);
+            } else {
+                config.output(file_name);
+            }
+
+            if let Some(steps) = &args.steps {
+                config.steps(*steps);
+            }
+
+            if args.random_seed {
+                config.seed(-1);
+            }
+
+            if let Some(width) = args.width {
+                config.width(width);
+            }
+
+            if let Some(height) = args.height {
+                config.height(height);
+            }
+
+            if let Some(negative) = args.negative {
+                config.negative_prompt(negative);
+            }
+
+            if args.low_vram {
+                model_config
+                    .clip_on_cpu(true)
+                    .vae_tiling(true)
+                    .flash_attention(true)
+                    .offload_params_to_cpu(true);
+            }
+
+            match args.preview {
+                Some(PreviewMode::Fast) => config.preview_mode(PreviewType::PREVIEW_PROJ),
+                Some(PreviewMode::Accurate) => config.preview_mode(PreviewType::PREVIEW_VAE),
+                None => config.preview_mode(PreviewType::PREVIEW_NONE),
+            };
+
+            if let Some(cache) = args.cache {
+                match cache {
+                    CacheMode::UCACHE => {
+                        config.ucache_caching(UCacheParamsBuilder::default().build().unwrap())
+                    }
+                    CacheMode::EASYCACHE => config
+                        .easy_cache_caching(EasyCacheParamsBuilder::default().build().unwrap()),
+                    CacheMode::DBCACHE => {
+                        config.db_cache_caching(DbCacheParamsBuilder::default().build().unwrap())
+                    }
+                    CacheMode::TAYLORSEER => config.taylor_seer_caching(),
+                    CacheMode::CACHEDIT => {
+                        config.cache_dit_caching(DbCacheParamsBuilder::default().build().unwrap())
+                    }
+                    CacheMode::SPECTRUM => config
+                        .spectrum_caching(SpectrumCacheParamsBuilder::default().build().unwrap()),
                 };
+            }
+            config.preview_output(preview_filename);
 
-                if let Some(cache) = args.cache {
-                    match cache {
-                        CacheMode::UCACHE => {
-                            config.ucache_caching(UCacheParamsBuilder::default().build().unwrap())
-                        }
-                        CacheMode::EASYCACHE => config
-                            .easy_cache_caching(EasyCacheParamsBuilder::default().build().unwrap()),
-                        CacheMode::DBCACHE => config
-                            .db_cache_caching(DbCacheParamsBuilder::default().build().unwrap()),
-                        CacheMode::TAYLORSEER => config.taylor_seer_caching(),
-                        CacheMode::CACHEDIT => config
-                            .cache_dit_caching(DbCacheParamsBuilder::default().build().unwrap()),
-                    };
-                }
-                config.preview_output(preview_filename);
-
-                Ok((config, model_config))
-            })
-            .build()
-            .unwrap();
+            Ok((config, model_config))
+        })
+        .build()
+        .unwrap();
     gen_img(&config, &mut model_config).unwrap();
 
     println!();

src/api.rs

@@ -117,6 +117,38 @@ pub struct LoraSpec {
     pub multiplier: f32,
 }
 
+/// Parameters for Spectrum Caching
+#[derive(Builder, Debug, Clone)]
+pub struct SpectrumCacheParams {
+    /// Chebyshev vs Taylor blend weight (0=Taylor, 1=Chebyshev)
+    #[builder(default = "0.40")]
+    w: f32,
+
+    /// Chebyshev polynomial degree
+    #[builder(default = "3")]
+    m: i32,
+
+    /// Ridge regression regularization
+    #[builder(default = "1.0")]
+    lam: f32,
+
+    /// Initial window size (compute every N steps)
+    #[builder(default = "2")]
+    window: i32,
+
+    /// Window growth per computed step after warmup
+    #[builder(default = "0.50")]
+    flex: f32,
+
+    /// Steps to always compute before caching starts
+    #[builder(default = "4")]
+    warmup: i32,
+
+    /// Stop caching at this fraction of total steps
+    #[builder(default = "0.9")]
+    stop: f32,
+}
+
 /// Parameters for UCache
 #[derive(Builder, Debug, Clone)]
 pub struct UCacheParams {
@@ -990,6 +1022,13 @@ impl ConfigBuilder {
             taylorseer_skip_interval: 1,
             scm_mask: null(),
             scm_policy_dynamic: true,
+            spectrum_w: 0.4,
+            spectrum_m: 3,
+            spectrum_lam: 1.0,
+            spectrum_window_size: 2,
+            spectrum_flex_window: 0.5,
+            spectrum_warmup_steps: 4,
+            spectrum_stop_percent: 0.9,
         }
     }
 
@@ -1000,6 +1039,20 @@ impl ConfigBuilder {
         self
     }
 
+    pub fn spectrum_caching(&mut self, params: SpectrumCacheParams) -> &mut Self {
+        let mut cache = Self::cache_init();
+        cache.mode = sd_cache_mode_t::SD_CACHE_SPECTRUM;
+        cache.spectrum_w = params.w;
+        cache.spectrum_m = params.m;
+        cache.spectrum_lam = params.lam;
+        cache.spectrum_window_size = params.window;
+        cache.spectrum_flex_window = params.flex;
+        cache.spectrum_warmup_steps = params.warmup;
+        cache.spectrum_stop_percent = params.stop;
+        self.cache = Some(cache);
+        self
+    }
+
     pub fn ucache_caching(&mut self, params: UCacheParams) -> &mut Self {
         let mut cache = Self::cache_init();
         cache.mode = sd_cache_mode_t::SD_CACHE_UCACHE;