From ae1cc2fac4d5d960410971620f3746762e673769 Mon Sep 17 00:00:00 2001 From: newfla Date: Tue, 16 Jun 2026 18:45:53 +0200 Subject: [PATCH] feat: support backend-specific max-vram budgets --- cli/README.md | 2 -- cli/src/main.rs | 6 ---- src/api.rs | 64 ++++++++++++++++++++++++++++++++++++---- sys/stable-diffusion.cpp | 2 +- 4 files changed, 60 insertions(+), 14 deletions(-) diff --git a/cli/README.md b/cli/README.md index 77cfc5c..7702b0c 100644 --- a/cli/README.md +++ b/cli/README.md @@ -39,8 +39,6 @@ Options: Set Huggingface Hub token. Only used when downloading models that have not been cached before --low-vram Enable optimization to use less VRAM: clip_on_cpu, vae tiling, flash_attention, offload_params_to_cpu - --max-vram - Maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables graph splitting; -1 auto-detects free VRAM minus 1 GiB [default: -1] -s, --seed RNG seed (-1 --> random) [default: -1] -h, --help diff --git a/cli/src/main.rs b/cli/src/main.rs index 5f2035d..e3f808a 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -119,10 +119,6 @@ struct Args { #[arg(long, default_value_t = false)] low_vram: bool, - /// Maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables graph splitting - #[arg(long, default_value_t = -1.0)] - max_vram: f32, - /// RNG seed (-1 --> random) #[arg(short, long, default_value_t = -1)] seed: i32, @@ -168,8 +164,6 @@ fn main() { config.seed(args.seed); - model_config.max_vram(args.max_vram); - if let Some(width) = args.width { config.width(width); } diff --git a/src/api.rs b/src/api.rs index 1cc15a2..10a8e6f 100644 --- a/src/api.rs +++ b/src/api.rs @@ -28,6 +28,7 @@ use diffusion_rs_sys::sd_img_gen_params_t; use diffusion_rs_sys::sd_img_gen_params_to_str; use diffusion_rs_sys::sd_lora_t; use diffusion_rs_sys::sd_pm_params_t; +use diffusion_rs_sys::sd_pulid_params_t; use diffusion_rs_sys::sd_sample_params_t; use diffusion_rs_sys::sd_set_preview_callback; use diffusion_rs_sys::sd_set_progress_callback; @@ -101,7 +102,7 @@ pub enum DiffusionError { } #[non_exhaustive] -#[derive(Clone, Debug, strum::Display, strum::EnumIter)] +#[derive(Clone, Debug, strum::Display, strum::EnumIter, PartialEq, Eq, Hash)] #[strum(serialize_all = "lowercase")] /// Backend devices pub enum BackendDevice { @@ -335,8 +336,8 @@ pub struct ModelConfig { enable_mmap: 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, + #[builder(default = "Self::default_max_vram()", setter(custom))] + max_vram: (HashMap, CLibString), /// Path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now #[builder(default = "Default::default()")] @@ -426,6 +427,14 @@ pub struct ModelConfig { #[builder(default = "Default::default()")] high_noise_diffusion_model: CLibPath, + /// Path to PULID weights. Identity is injected during the denoise loop when paired with pulid_id_embedding. + #[builder(default = "Default::default()")] + pulid_weights_path: CLibPath, + + /// path to a .pulidembd binary produced by pulid_extract_id.py. Carries a (32, 2048) identity embedding extracted from a source portrait. Pair with pulid_weights on the context. + #[builder(default = "Default::default()")] + pulid_id_embedding_path: CLibPath, + /// Process vae in tiles to reduce memory usage (default: false) #[builder(default = "false")] vae_tiling: bool, @@ -719,6 +728,45 @@ impl ModelConfigBuilder { self } + pub fn max_vram(&mut self, vram_map: HashMap) -> &mut Self { + let vram_str = Self::concatenate_max_vram(&vram_map); + self.max_vram = Some((vram_map, vram_str)); + self + } + + fn default_max_vram() -> (HashMap, CLibString) { + let mut map = HashMap::new(); + #[cfg(feature = "metal")] + { + map.insert(BackendDevice::METAL, -1); + } + + #[cfg(feature = "vulkan")] + { + map.insert(BackendDevice::VULKAN0, -1); + } + + #[cfg(feature = "cuda")] + { + map.insert(BackendDevice::CUDA0, -1); + } + + map.insert(BackendDevice::CPU, -1); + + let map_str = Self::concatenate_max_vram(&map); + (map, map_str) + } + + fn concatenate_max_vram(vram_map: &HashMap) -> CLibString { + CLibString::from( + vram_map + .iter() + .map(|(key, value)| format!("{}={}", key, value)) + .collect::>() + .join(","), + ) + } + pub fn extra_tiling_args( &mut self, extra_tiling_args_map: HashMap, @@ -806,13 +854,15 @@ impl ModelConfig { circular_y: self.circular || self.circular_y, qwen_image_zero_cond_t: self.use_qwen_image_zero_cond_true, enable_mmap: self.enable_mmap, - max_vram: self.max_vram, + max_vram: self.max_vram.1.as_ptr(), backend: self.backend.1.as_ptr(), params_backend: self.params_backend.1.as_ptr(), embeddings_connectors_path: self.embeddings_connectors.as_ptr(), audio_vae_path: self.audio_vae.as_ptr(), vae_format: self.vae_format, stream_layers: self.stream_layers, + rpc_servers: null(), + pulid_weights_path: self.pulid_weights_path.as_ptr(), }; let ctx = new_sd_ctx(&sd_ctx_params); self.diffusion_ctx = Some((ctx, sd_ctx_params)) @@ -847,7 +897,7 @@ impl From<&ModelConfig> for ModelConfigBuilder { .map(|f| PathBuf::from_str(&f.0.into_string().unwrap()).unwrap()); builder .n_threads(value.n_threads) - .max_vram(value.max_vram) + .max_vram(value.max_vram.0.clone()) .upscale_repeats(value.upscale_repeats) .model(value.model.clone()) .diffusion_model(value.diffusion_model.clone()) @@ -1569,6 +1619,10 @@ fn gen_img_maybe_progress( loras: loras.as_ptr(), lora_count: loras.len() as u32, hires, + pulid_params: sd_pulid_params_t { + id_embedding_path: model_config.pulid_id_embedding_path.as_ptr(), + id_weight: 1.0, + }, }; let params_str = CString::from_raw(sd_img_gen_params_to_str(&sd_img_gen_params)) diff --git a/sys/stable-diffusion.cpp b/sys/stable-diffusion.cpp index 9838264..5a34bc7 160000 --- a/sys/stable-diffusion.cpp +++ b/sys/stable-diffusion.cpp @@ -1 +1 @@ -Subproject commit 9838264c490e2f413ed36b76aba9091b78ba4414 +Subproject commit 5a34bc7f6e0621dd2f899daa64476eac667d7ed3