Feat_flux2_klein_preset (#76)

* feat: sdxs preset

* fix(doc): improve help output

* feat: flux 2 klein presets

* fix: qwen3 9B download link

* fix(cli): enable flash attention when low-vram flag is set

Author: newfla

cli/README.md

@@ -2,28 +2,28 @@
 
 ## Usage
 ```
-cli.exe --preset flux1schnell --random-seed "a lovely duck holding a sign says 'drink your water'"
+cli.exe flux1schnell  "a lovely duck holding a sign says 'drink your water'" --random-seed --token YourHfHubToken
 ```
 ## Help
 ```
 Usage: cli.exe [OPTIONS] --preset <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]
   <PROMPT>  The prompt to render
 
 Options:
   -n, --negative <NEGATIVE>  Negative prompt
-  -p, --preset <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]
-  -w, --weights <WEIGHTS>    Optionally which type of quantization to use [possible values: F32, F16, Q4_0, Q4_1, Q5_0, Q5_1, Q8_0, Q8_1, Q2_K, Q3_K, Q4_K, Q5_K, Q6_K, Q8_K, IQ2_XXS, IQ2_XS, IQ3_XXS, IQ1_S, IQ4_NL, IQ3_S, IQ2_S, IQ4_XS, I8, I16, I32, I64, F64, IQ1_M, BF16, TQ1_0, TQ2_0, MXFP4, F8_E4M3]
-  -s, --steps <STEPS>        Numer of inference steps
-  -w, --width <WIDTH>        Width
-  -h, --height <HEIGHT>      Height
-  -b, --batch <BATCH>        Number of images to generate [default: 1]
+      --weights <WEIGHTS>    Optionally which type of quantization to use [possible values: F32, F16, Q4_0, Q4_1, Q5_0, Q5_1, Q8_0, Q8_1, Q2_K, Q3_K, Q4_K, Q5_K, Q6_K, Q8_K, IQ2_XXS, IQ2_XS, IQ3_XXS, IQ1_S, IQ4_NL, IQ3_S, IQ2_S, IQ4_XS, I8, I16, I32, I64, F64, IQ1_M, BF16, TQ1_0, TQ2_0, MXFP4, F8_E4M3]
+  -s, --steps <STEPS>        Override the preset default number of inference steps
+      --width <WIDTH>        Override the preset default width
+      --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: ./]
   -p, --preview <PREVIEW>    Enable preview [possible values: fast, accurate]
-  -t, --token <TOKEN>        Set Huggingface Hub token
-  -l, --low-vram             Enable optimization for gpu with lower GB
-  -r, --random-seed          Enable Random Seed: different runs will produce different results
+  -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      
   -h, --help                 Print help
   -V, --version              Print version
   ```

cli/src/main.rs

@@ -7,6 +7,7 @@ use diffusion_rs::{
     api::{PreviewType, gen_img},
     preset::{
         ChromaRadianceWeight, ChromaWeight, DiffInstructStarWeight, Flux1MiniWeight, Flux1Weight,
+        Flux2Klein4BWeight, Flux2Klein9BWeight, Flux2KleinBase4BWeight, Flux2KleinBase9BWeight,
         Flux2Weight, NitroSDRealismWeight, NitroSDVibrantWeight, OvisImageWeight, Preset,
         PresetBuilder, PresetDiscriminants, QwenImageWeight, SSD1BWeight,
         TwinFlowZImageTurboExpWeight, WeightType, ZImageTurboWeight,
@@ -32,34 +33,34 @@ enum PreviewMode {
 #[derive(Parser)]
 #[command(version, about, long_about = None)]
 struct Args {
+    /// The preset to use
+    #[arg(ignore_case = true, value_parser = clap_enum_variants!(PresetDiscriminants))]
+    preset: PresetDiscriminants,
+
     /// The prompt to render
     prompt: String,
 
     /// Negative prompt
     #[arg(short, long)]
     negative: Option<String>,
 
-    /// The preset to use
-    #[arg(short, long, ignore_case = true, value_parser = clap_enum_variants!(PresetDiscriminants))]
-    preset: PresetDiscriminants,
-
     /// Optionally which type of quantization to use
-    #[arg(short, long, ignore_case = true, value_parser = clap_enum_variants!(WeightType))]
+    #[arg(long, ignore_case = true, value_parser = clap_enum_variants!(WeightType))]
     weights: Option<WeightType>,
 
-    /// Numer of inference steps
+    /// Override the preset default number of inference steps
     #[arg(short, long)]
     steps: Option<i32>,
 
-    /// Width
-    #[arg(short, long)]
+    /// Override the preset default width
+    #[arg(long)]
     width: Option<i32>,
 
-    /// Height
-    #[arg(short, long)]
+    /// Override the preset default height
+    #[arg(long)]
     height: Option<i32>,
 
-    /// Number of images to generate
+    /// Number of images to generate (default 1)
     #[arg(short, long, default_value_t = 1)]
     batch: i32,
 
@@ -71,11 +72,11 @@ struct Args {
     #[arg(short, long, ignore_case = true)]
     preview: Option<PreviewMode>,
 
-    /// Set Huggingface Hub token
+    /// Set Huggingface Hub token. Only used when downloading models that have not been cached before
     #[arg(short, long)]
     token: Option<String>,
 
-    /// Enable optimization for gpu with lower GB
+    /// Enable optimization to use less VRAM: clip_on_cpu, vae tiling, flash_attention, offload_params_to_cpu
     #[arg(short, long, default_value_t = false)]
     low_vram: bool,
 
@@ -142,6 +143,7 @@ fn main() {
                 model_config
                     .clip_on_cpu(true)
                     .vae_tiling(true)
+                    .flash_attention(true)
                     .offload_params_to_cpu(true);
             }
 
@@ -268,6 +270,31 @@ fn get_preset(args: &Args) -> Preset {
                 .try_into()
                 .unwrap(),
         ),
+        PresetDiscriminants::SDXS512DreamShaper => Preset::SDXS512DreamShaper,
+        PresetDiscriminants::Flux2Klein4B => Preset::Flux2Klein4B(
+            args.weights
+                .unwrap_or_else(|| Flux2Klein4BWeight::default().into())
+                .try_into()
+                .unwrap(),
+        ),
+        PresetDiscriminants::Flux2KleinBase4B => Preset::Flux2KleinBase4B(
+            args.weights
+                .unwrap_or_else(|| Flux2KleinBase4BWeight::default().into())
+                .try_into()
+                .unwrap(),
+        ),
+        PresetDiscriminants::Flux2Klein9B => Preset::Flux2Klein9B(
+            args.weights
+                .unwrap_or_else(|| Flux2Klein9BWeight::default().into())
+                .try_into()
+                .unwrap(),
+        ),
+        PresetDiscriminants::Flux2KleinBase9B => Preset::Flux2KleinBase9B(
+            args.weights
+                .unwrap_or_else(|| Flux2KleinBase9BWeight::default().into())
+                .try_into()
+                .unwrap(),
+        ),
     };
     preset
 }

src/preset.rs

@@ -7,11 +7,13 @@ use crate::{
     api::{Config, ConfigBuilder, ConfigBuilderError, ModelConfig, ModelConfigBuilder},
     preset_builder::{
         chroma, chroma_radiance, diff_instruct_star, dream_shaper_xl_2_1_turbo, flux_1_dev,
-        flux_1_mini, flux_1_schnell, flux_2_dev, juggernaut_xl_11, nitro_sd_realism,
-        nitro_sd_vibrant, ovis_image, qwen_image, sd_turbo, sdxl_base_1_0, sdxl_turbo_1_0, ssd_1b,
-        stable_diffusion_1_4, stable_diffusion_1_5, stable_diffusion_2_1,
-        stable_diffusion_3_5_large, stable_diffusion_3_5_large_turbo, stable_diffusion_3_5_medium,
-        stable_diffusion_3_medium, twinflow_z_image_turbo, z_image_turbo,
+        flux_1_mini, flux_1_schnell, flux_2_dev, flux_2_klein_4b, flux_2_klein_9b,
+        flux_2_klein_base_4b, flux_2_klein_base_9b, juggernaut_xl_11, nitro_sd_realism,
+        nitro_sd_vibrant, ovis_image, qwen_image, sd_turbo, sdxl_base_1_0, sdxl_turbo_1_0,
+        sdxs512_dream_shaper, ssd_1b, stable_diffusion_1_4, stable_diffusion_1_5,
+        stable_diffusion_2_1, stable_diffusion_3_5_large, stable_diffusion_3_5_large_turbo,
+        stable_diffusion_3_5_medium, stable_diffusion_3_medium, twinflow_z_image_turbo,
+        z_image_turbo,
     },
 };
 
@@ -30,7 +32,11 @@ use crate::{
     ZImageTurboWeight(derive(Default)),
     QwenImageWeight(derive(Default)),
     OvisImageWeight(derive(Default)),
-    TwinFlowZImageTurboExpWeight(derive(Default))
+    TwinFlowZImageTurboExpWeight(derive(Default)),
+    Flux2Klein4BWeight(derive(Default)),
+    Flux2KleinBase4BWeight(derive(Default)),
+    Flux2Klein9BWeight(derive(Default)),
+    Flux2KleinBase9BWeight(derive(Default))
 )]
 #[derive(Debug, Clone, Copy, EnumString, VariantNames)]
 #[strum(ascii_case_insensitive)]
@@ -55,7 +61,11 @@ pub enum WeightType {
         ZImageTurboWeight,
         QwenImageWeight,
         OvisImageWeight(default),
-        TwinFlowZImageTurboExpWeight(default)
+        TwinFlowZImageTurboExpWeight(default),
+        Flux2Klein4BWeight,
+        Flux2KleinBase4BWeight,
+        Flux2Klein9BWeight(default),
+        Flux2KleinBase9BWeight(default)
     )]
     Q4_0,
     #[subenum(Flux2Weight, QwenImageWeight)]
@@ -84,7 +94,10 @@ pub enum WeightType {
         ZImageTurboWeight,
         QwenImageWeight,
         OvisImageWeight,
-        TwinFlowZImageTurboExpWeight
+        TwinFlowZImageTurboExpWeight,
+        Flux2Klein4BWeight(default),
+        Flux2KleinBase4BWeight(default),
+        Flux2Klein9BWeight
     )]
     Q8_0,
     Q8_1,
@@ -149,7 +162,11 @@ pub enum WeightType {
         ZImageTurboWeight,
         QwenImageWeight,
         OvisImageWeight,
-        TwinFlowZImageTurboExpWeight
+        TwinFlowZImageTurboExpWeight,
+        Flux2Klein4BWeight,
+        Flux2KleinBase4BWeight,
+        Flux2Klein9BWeight,
+        Flux2KleinBase9BWeight
     )]
     BF16,
     TQ1_0,
@@ -229,6 +246,20 @@ pub enum Preset {
     /// Requires access rights to <https://huggingface.co/black-forest-labs/FLUX.1-schnell> providing a token via [crate::util::set_hf_token]
     /// Enabled [crate::api::SampleMethod::DPM2_SAMPLE_METHOD] and [crate::api::Scheduler::SMOOTHSTEP_SCHEDULER]. cfg_scale 1.0. 3 steps. Flash attention enabled. 1024x512. Vae-tiling enabled.
     TwinFlowZImageTurboExp(TwinFlowZImageTurboExpWeight),
+    /// cfg_scale 1.0. 1 steps. 512x512
+    SDXS512DreamShaper,
+    /// Requires access rights to <https://huggingface.co/black-forest-labs/FLUX.2-dev> providing a token via [crate::util::set_hf_token]
+    /// cfg scale 1.0. 4 steps. Flash attention enabled. Offload params to CPU enabled. 1024x1024. Vae-tiling enabled
+    Flux2Klein4B(Flux2Klein4BWeight),
+    /// Requires access rights to <https://huggingface.co/black-forest-labs/FLUX.2-dev> providing a token via [crate::util::set_hf_token]
+    /// cfg scale 4.0. 20 steps. Flash attention enabled. Offload params to CPU enabled. 1024x1024. Vae-tiling enabled
+    Flux2KleinBase4B(Flux2KleinBase4BWeight),
+    /// Requires access rights to <https://huggingface.co/black-forest-labs/FLUX.2-dev> providing a token via [crate::util::set_hf_token]
+    /// cfg scale 1.0. 4 steps. Flash attention enabled. Offload params to CPU enabled. 1024x1024. Vae-tiling enabled
+    Flux2Klein9B(Flux2Klein9BWeight),
+    /// Requires access rights to <https://huggingface.co/black-forest-labs/FLUX.2-dev> providing a token via [crate::util::set_hf_token]
+    /// cfg scale 4.0. 20 steps. Flash attention enabled. Offload params to CPU enabled. 1024x1024. Vae-tiling enabled
+    Flux2KleinBase9B(Flux2KleinBase9BWeight),
 }
 
 impl Preset {
@@ -261,6 +292,11 @@ impl Preset {
             Preset::OvisImage(sd_type_t) => ovis_image(sd_type_t),
             Preset::DreamShaperXL2_1Turbo => dream_shaper_xl_2_1_turbo(),
             Preset::TwinFlowZImageTurboExp(sd_type_t) => twinflow_z_image_turbo(sd_type_t),
+            Preset::SDXS512DreamShaper => sdxs512_dream_shaper(),
+            Preset::Flux2Klein4B(sd_type_t) => flux_2_klein_4b(sd_type_t),
+            Preset::Flux2KleinBase4B(sd_type_t) => flux_2_klein_base_4b(sd_type_t),
+            Preset::Flux2Klein9B(sd_type_t) => flux_2_klein_9b(sd_type_t),
+            Preset::Flux2KleinBase9B(sd_type_t) => flux_2_klein_base_9b(sd_type_t),
         };
 
         // Metal workaround.
@@ -343,8 +379,10 @@ mod tests {
         api::gen_img,
         preset::{
             ChromaRadianceWeight, ChromaWeight, DiffInstructStarWeight, Flux1MiniWeight,
-            Flux1Weight, Flux2Weight, NitroSDRealismWeight, NitroSDVibrantWeight, OvisImageWeight,
-            QwenImageWeight, SSD1BWeight, TwinFlowZImageTurboExpWeight, ZImageTurboWeight,
+            Flux1Weight, Flux2Klein4BWeight, Flux2Klein9BWeight, Flux2KleinBase4BWeight,
+            Flux2KleinBase9BWeight, Flux2Weight, NitroSDRealismWeight, NitroSDVibrantWeight,
+            OvisImageWeight, QwenImageWeight, SSD1BWeight, TwinFlowZImageTurboExpWeight,
+            ZImageTurboWeight,
         },
         util::set_hf_token,
     };
@@ -531,4 +569,38 @@ mod tests {
             TwinFlowZImageTurboExpWeight::Q3_K,
         ));
     }
+
+    #[ignore]
+    #[test]
+    fn test_sdxs512_dream_shaper() {
+        run(Preset::SDXS512DreamShaper);
+    }
+
+    #[ignore]
+    #[test]
+    fn test_flux_2_klein_4b() {
+        set_hf_token(include_str!("../token.txt"));
+        run(Preset::Flux2Klein4B(Flux2Klein4BWeight::Q8_0));
+    }
+
+    #[ignore]
+    #[test]
+    fn test_flux_2_klein_base_4b() {
+        set_hf_token(include_str!("../token.txt"));
+        run(Preset::Flux2KleinBase4B(Flux2KleinBase4BWeight::Q8_0));
+    }
+
+    #[ignore]
+    #[test]
+    fn test_flux_2_klein_9b() {
+        set_hf_token(include_str!("../token.txt"));
+        run(Preset::Flux2Klein9B(Flux2Klein9BWeight::Q4_0));
+    }
+
+    #[ignore]
+    #[test]
+    fn test_flux_2_klein_base_9b() {
+        set_hf_token(include_str!("../token.txt"));
+        run(Preset::Flux2KleinBase9B(Flux2KleinBase9BWeight::Q4_0));
+    }
 }