From c8fa722b3deef45062ab5f5126ecda3f204c0459 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 2 Feb 2024 23:17:12 -0500 Subject: [PATCH] parameter for predefined SD weights --- .../src/job/job_types/image_generation/sd/process_job.rs | 5 ++++- .../image_generation/sd/stable_diffusion_dependencies.rs | 4 ++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs index 050e3b2320..8498767dd5 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs @@ -416,7 +416,10 @@ pub async fn process_job_lora( // download model + vae for lora // use lora with whatever checkpoint make it fixed id if it works submit it - let weight_token_sd_model_token = ModelWeightToken(String::from("weight_dmmthavhawqc2hj7yqyemcbf8")); // any generic model. + let weight_token_sd_model_token = ModelWeightToken::new_from_str(&sd_deps.predefined_sd_weight_token); + + info!("Using predefined SD weight token: {:?}", &weight_token_sd_model_token); + let sd_weight_record = get_weight_by_token( &weight_token_sd_model_token, false, diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs index 79d7e69d8a..f89c15a9ac 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs @@ -8,6 +8,8 @@ pub struct StableDiffusionDependencies { pub downloaders: StableDiffusionDownloaders, pub inference_command: StableDiffusionInferenceCommand, pub vae_bucket_path: String, + pub predefined_sd_weight_token: String, + pub predefined_lora_weight_token: String, } impl StableDiffusionDependencies { @@ -16,6 +18,8 @@ impl StableDiffusionDependencies { downloaders: StableDiffusionDownloaders::build_all_from_env(), inference_command: StableDiffusionInferenceCommand::from_env()?, vae_bucket_path: easyenv::get_env_string_required("SD_VAE_BUCKET_PATH")?, + predefined_sd_weight_token: easyenv::get_env_string_required("SD_PREDEFINED_SD_WEIGHT_TOKEN")?, + predefined_lora_weight_token: easyenv::get_env_string_required("SD_PREDEFINED_LORA_WEIGHT_TOKEN")?, }) } }