let model be configurable
This commit is contained in:
parent
b5d8842ccc
commit
22f79bbdf1
1 changed files with 10 additions and 1 deletions
11
src/main.rs
11
src/main.rs
|
@ -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);
|
||||
|
||||
|
|
Loading…
Reference in a new issue