diffusion_rs/
modifier.rs

1use hf_hub::api::sync::ApiError;
2
3use crate::{api::SampleMethod, preset::ConfigsBuilder, util::download_file_hf_hub};
4
5/// Add the <https://huggingface.co/ximso/RealESRGAN_x4plus_anime_6B> upscaler
6pub fn real_esrgan_x4plus_anime_6_b(
7    mut builder: ConfigsBuilder,
8) -> Result<ConfigsBuilder, ApiError> {
9    let upscaler_path = download_file_hf_hub(
10        "ximso/RealESRGAN_x4plus_anime_6B",
11        "RealESRGAN_x4plus_anime_6B.pth",
12    )?;
13    builder.1.upscale_model(upscaler_path);
14    Ok(builder)
15}
16
17/// Apply <https://huggingface.co/madebyollin/sdxl-vae-fp16-fix> to avoid black images with xl models
18pub fn sdxl_vae_fp16_fix(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
19    let vae_path = download_file_hf_hub("madebyollin/sdxl-vae-fp16-fix", "sdxl.vae.safetensors")?;
20    builder.1.vae(vae_path);
21    Ok(builder)
22}
23
24/// Apply <https://huggingface.co/madebyollin/taesd> taesd autoencoder for faster decoding (SD v1/v2)
25pub fn taesd(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
26    let taesd_path =
27        download_file_hf_hub("madebyollin/taesd", "diffusion_pytorch_model.safetensors")?;
28    builder.1.taesd(taesd_path);
29    Ok(builder)
30}
31
32/// Apply <https://huggingface.co/madebyollin/taesdxl> taesd autoencoder for faster decoding (SDXL)
33pub fn taesd_xl(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
34    let taesd_path =
35        download_file_hf_hub("madebyollin/taesdxl", "diffusion_pytorch_model.safetensors")?;
36    builder.1.taesd(taesd_path);
37    Ok(builder)
38}
39
40/// Apply <https://huggingface.co/cqyan/hybrid-sd-tinyvae> taesd autoencoder for faster decoding (SD v1/v2)
41pub fn hybrid_taesd(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
42    let taesd_path = download_file_hf_hub(
43        "cqyan/hybrid-sd-tinyvae",
44        "diffusion_pytorch_model.safetensors",
45    )?;
46    builder.1.taesd(taesd_path);
47    Ok(builder)
48}
49
50/// Apply <https://huggingface.co/cqyan/hybrid-sd-tinyvae-xl> taesd autoencoder for faster decoding (SDXL)
51pub fn hybrid_taesd_xl(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
52    let taesd_path = download_file_hf_hub(
53        "cqyan/hybrid-sd-tinyvae-xl",
54        "diffusion_pytorch_model.safetensors",
55    )?;
56    builder.1.taesd(taesd_path);
57    Ok(builder)
58}
59
60/// Apply <https://huggingface.co/latent-consistency/lcm-lora-sdv1-5> to reduce inference steps for SD v1 between 2-8
61/// cfg_scale 1. 4 steps.
62pub fn lcm_lora_sd_1_5(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
63    let lora_path = download_file_hf_hub(
64        "latent-consistency/lcm-lora-sdv1-5",
65        "pytorch_lora_weights.safetensors",
66    )?;
67    builder.1.lora_model(&lora_path);
68    builder.0.cfg_scale(1.).steps(4);
69    Ok(builder)
70}
71
72/// Apply <https://huggingface.co/latent-consistency/lcm-lora-sdxl> to reduce inference steps for SD v1 between 2-8 (default 8)
73/// Enabled [SampleMethod::LCM]. cfg_scale 2. 8 steps.
74pub fn lcm_lora_sdxl_base_1_0(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
75    let lora_path = download_file_hf_hub(
76        "latent-consistency/lcm-lora-sdxl",
77        "pytorch_lora_weights.safetensors",
78    )?;
79    builder.1.lora_model(&lora_path);
80    builder
81        .0
82        .cfg_scale(2.)
83        .steps(8)
84        .sampling_method(SampleMethod::LCM);
85    Ok(builder)
86}
87
88/// Apply <https://huggingface.co/comfyanonymous/flux_text_encoders/blob/main/t5xxl_fp8_e4m3fn.safetensors> Fp8 t5xxl text encoder to reduce memory usage
89pub fn t5xxl_fp8_flux_1(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
90    let t5xxl_path = download_file_hf_hub(
91        "comfyanonymous/flux_text_encoders",
92        "t5xxl_fp8_e4m3fn.safetensors",
93    )?;
94
95    builder.1.t5xxl(t5xxl_path);
96    Ok(builder)
97}
98
99/// Apply <https://huggingface.co/comfyanonymous/flux_text_encoders/blob/main/t5xxl_fp16.safetensors>
100/// Default for flux_1_dev/schnell
101pub fn t5xxl_fp16_flux_1(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
102    let t5xxl_path = download_file_hf_hub(
103        "comfyanonymous/flux_text_encoders",
104        "t5xxl_fp16.safetensors",
105    )?;
106
107    builder.1.t5xxl(t5xxl_path);
108    Ok(builder)
109}
110
111/// Apply <https://huggingface.co/Green-Sky/flux.1-schnell-GGUF/blob/main/t5xxl_q2_k.gguf>
112pub fn t5xxl_q2_k_flux_1(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
113    let t5xxl_path = download_file_hf_hub("Green-Sky/flux.1-schnell-GGUF", "t5xxl_q2_k.gguf")?;
114
115    builder.1.t5xxl(t5xxl_path);
116    Ok(builder)
117}
118
119/// Apply <https://huggingface.co/Green-Sky/flux.1-schnell-GGUF/blob/main/t5xxl_q3_k.gguf>
120pub fn t5xxl_q3_k_flux_1(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
121    let t5xxl_path = download_file_hf_hub("Green-Sky/flux.1-schnell-GGUF", "t5xxl_q3_k.gguf")?;
122
123    builder.1.t5xxl(t5xxl_path);
124    Ok(builder)
125}
126
127/// Apply <https://huggingface.co/Green-Sky/flux.1-schnell-GGUF/blob/main/t5xxl_q4_k.gguf>
128/// Default for flux_1_mini
129pub fn t5xxl_q4_k_flux_1(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
130    let t5xxl_path = download_file_hf_hub("Green-Sky/flux.1-schnell-GGUF", "t5xxl_q4_k.gguf")?;
131
132    builder.1.t5xxl(t5xxl_path);
133    Ok(builder)
134}
135
136/// Apply <https://huggingface.co/Green-Sky/flux.1-schnell-GGUF/blob/main/t5xxl_q8_0.gguf>
137pub fn t5xxl_q8_0_flux_1(mut builder: ConfigsBuilder) -> Result<ConfigsBuilder, ApiError> {
138    let t5xxl_path = download_file_hf_hub("Green-Sky/flux.1-schnell-GGUF", "t5xxl_q8_0.gguf")?;
139
140    builder.1.t5xxl(t5xxl_path);
141    Ok(builder)
142}
143
144#[cfg(test)]
145mod tests {
146    use crate::{
147        api::txt2img,
148        preset::{Modifier, Preset, PresetBuilder},
149    };
150
151    use super::{
152        hybrid_taesd, hybrid_taesd_xl, lcm_lora_sd_1_5, lcm_lora_sdxl_base_1_0, taesd, taesd_xl,
153    };
154
155    static PROMPT: &str = "a lovely duck drinking water from a bottle";
156
157    fn run(preset: Preset, m: Modifier) {
158        let (mut config, mut model_config) = PresetBuilder::default()
159            .preset(preset)
160            .prompt(PROMPT)
161            .with_modifier(m)
162            .build()
163            .unwrap();
164        txt2img(&mut config, &mut model_config).unwrap();
165    }
166
167    #[ignore]
168    #[test]
169    fn test_taesd() {
170        run(Preset::StableDiffusion1_5, taesd);
171    }
172
173    #[ignore]
174    #[test]
175    fn test_taesd_xl() {
176        run(Preset::SDXLTurbo1_0Fp16, taesd_xl);
177    }
178
179    #[ignore]
180    #[test]
181    fn test_hybrid_taesd() {
182        run(Preset::StableDiffusion1_5, hybrid_taesd);
183    }
184
185    #[ignore]
186    #[test]
187    fn test_hybrid_taesd_xl() {
188        run(Preset::SDXLTurbo1_0Fp16, hybrid_taesd_xl);
189    }
190
191    #[ignore]
192    #[test]
193    fn test_lcm_lora_sd_1_5() {
194        run(Preset::StableDiffusion1_5, lcm_lora_sd_1_5);
195    }
196
197    #[ignore]
198    #[test]
199    fn test_lcm_lora_sdxl_base_1_0() {
200        run(Preset::SDXLBase1_0, lcm_lora_sdxl_base_1_0);
201    }
202}
diffusion_rs/modifier.rs

diffusion_rs/
modifier.rs
