let model be configurable

This commit is contained in:
Luna 2024-08-04 18:35:21 -03:00
parent b5d8842ccc
commit 22f79bbdf1

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]
async fn main() {
pretty_env_logger::init();
@ -130,7 +137,9 @@ async fn send_image_to_dd(
let mut map: HashMap<&str, &str> = HashMap::new();
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("image", &file_base64);