luna/glimbus
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

let model be configurable

Browse source
This commit is contained in:
Luna 2024-08-04 18:35:21 -03:00
parent b5d8842ccc
commit 22f79bbdf1
1 changed files with 10 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
src/main.rs
View file

@ -31,6 +31,13 @@ fn get_upstream_runner() -> Runner {
} }
} }
fn get_sd_model() -> String {
match std::env::var("SD_MODEL") {
Ok(value) => value,
Err(_) => "wd14-vit-v2-git".to_string(),
}
}
#[tokio::main] #[tokio::main]
async fn main() { async fn main() {
pretty_env_logger::init(); pretty_env_logger::init();
@ -130,7 +137,9 @@ async fn send_image_to_dd(
let mut map: HashMap<&str, &str> = HashMap::new(); let mut map: HashMap<&str, &str> = HashMap::new();
let file_base64 = general_purpose::STANDARD.encode(file_contents.clone()); let file_base64 = general_purpose::STANDARD.encode(file_contents.clone());
map.insert("model", "wd14-vit-v2-git"); let sd_model = get_sd_model();
map.insert("model", &sd_model);
map.insert("threshold", options.threshold.as_ref()); map.insert("threshold", options.threshold.as_ref());
map.insert("image", &file_base64); map.insert("image", &file_base64);