daniel543x/artcraft
1
0 Fork 0
mirror of https://github.com/storytold/artcraft.git synced 2026-11-06 02:00:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

hack for VAE SD support

Browse Source
This commit is contained in:
Brandon Thomas committed 2024-02-02 22:58:28 -05:00
1 parent 449c29e824
commit 0540d2caba
2 files changed
+110 -74
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs
+108 -74
View File
@@ -3,7 +3,7 @@ use std::time::{Duration, Instant};
use actix_web::dev::ResourcePath;
use anyhow::anyhow;
use log::info;
use log::{error, info};
use serde_json;
use cloud_storage::remote_file_manager::media_descriptor::MediaImagePngDescriptor;
@@ -75,6 +75,7 @@ pub async fn sd_args_from_job(
| args.args.as_ref()
)
.flatten();
let polymorphic_args = match inference_args {
Some(args) => args,
None => {
@@ -92,6 +93,7 @@ pub async fn sd_args_from_job(
);
}
};
let stable_diffusion_args: StableDiffusionArgs = StableDiffusionArgs::from(sd_args.clone());
Ok(stable_diffusion_args)
}
@@ -208,32 +210,44 @@ pub async fn process_job_sd(
// download vae for model
// use this vae doesn't matter though
// VAE token for now
let vae_token = String::from("weight_rb0959wfzjhk3d1k93hr3s0qw");
let model_weight_vae = ModelWeightToken(vae_token);
let vae_weight_record = get_weight_by_token(
&model_weight_vae,
false,
&deps.db.mysql_pool
).await?;
let vae_weight_record = match vae_weight_record {
Some(val) => val,
None => {
return Err(
ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))
);
}
};
args.job_dependencies
.buckets
.public_bucket_client
.download_file_to_disk(&sd_deps.vae_bucket_path, &vae_path)
.await
.map_err(|err| {
error!("could not download VAE: {:?}", err);
ProcessSingleJobError::from_anyhow_error(anyhow!("could not download VAE: {:?}", err))
})?;
let vae_details = RemoteCloudBucketDetails::new(
vae_weight_record.public_bucket_hash.clone(),
vae_weight_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()),
vae_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string())
);
// // use this vae doesn't matter though
// // VAE token for now
// let vae_token = String::from("REPLACE_ME");
// let model_weight_vae = ModelWeightToken(vae_token);
// let vae_weight_record = get_weight_by_token(
// &model_weight_vae,
// false,
// &deps.db.mysql_pool
// ).await?;
// let vae_weight_record = match vae_weight_record {
// Some(val) => val,
// None => {
// return Err(
// ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))
// );
// }
// };
//
// let vae_details = RemoteCloudBucketDetails::new(
// vae_weight_record.public_bucket_hash.clone(),
// vae_weight_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()),
// vae_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string())
// );
//
// let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?;
// remote_cloud_file_client.download_file(vae_details, path_to_string(vae_path.clone())).await?;
let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?;
remote_cloud_file_client.download_file(vae_details, path_to_string(vae_path.clone())).await?;
let stderr_output_file = work_temp_dir.path().join("sd_err.txt");
let stdout_output_file = work_temp_dir.path().join("sd_out.txt");
@@ -427,32 +441,42 @@ pub async fn process_job_lora(
remote_cloud_file_client.download_file(sd_weight_details, path_to_string(sd_checkpoint_path.clone())).await?;
// use this vae doesn't matter though
// VAE token for now
let vae_token = String::from("weight_rb0959wfzjhk3d1k93hr3s0qw");
let model_weight_vae = ModelWeightToken(vae_token);
let vae_weight_record = get_weight_by_token(
&model_weight_vae,
false,
&deps.db.mysql_pool
).await?;
let vae_weight_record = match vae_weight_record {
Some(val) => val,
None => {
return Err(
ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))
);
}
};
args.job_dependencies
.buckets
.public_bucket_client
.download_file_to_disk(&sd_deps.vae_bucket_path, &vae_path)
.await
.map_err(|err| {
error!("could not download VAE: {:?}", err);
ProcessSingleJobError::from_anyhow_error(anyhow!("could not download VAE: {:?}", err))
})?;
let vae_details = RemoteCloudBucketDetails::new(
vae_weight_record.public_bucket_hash.clone(),
vae_weight_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()),
vae_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string())
);
let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?;
remote_cloud_file_client.download_file(vae_details, path_to_string(vae_path.clone())).await?;
// // use this vae doesn't matter though
// // VAE token for now
// let vae_token = String::from("REPLACE_ME");
// let model_weight_vae = ModelWeightToken(vae_token);
// let vae_weight_record = get_weight_by_token(
// &model_weight_vae,
// false,
// &deps.db.mysql_pool
// ).await?;
// let vae_weight_record = match vae_weight_record {
// Some(val) => val,
// None => {
// return Err(
// ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))
// );
// }
// };
//
// let vae_details = RemoteCloudBucketDetails::new(
// vae_weight_record.public_bucket_hash.clone(),
// vae_weight_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()),
// vae_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string())
// );
//
// let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?;
// remote_cloud_file_client.download_file(vae_details, path_to_string(vae_path.clone())).await?;
let stderr_output_file = work_temp_dir.path().join("sd_err.txt");
let stdout_output_file = work_temp_dir.path().join("sd_out.txt");
@@ -688,32 +712,42 @@ pub async fn process_job_inference(
}
}
// VAE token for now
let vae_token = String::from("weight_rb0959wfzjhk3d1k93hr3s0qw");
let model_weight_vae = ModelWeightToken(vae_token);
args.job_dependencies
.buckets
.public_bucket_client
.download_file_to_disk(&sd_deps.vae_bucket_path, &vae_path)
.await
.map_err(|err| {
error!("could not download VAE: {:?}", err);
ProcessSingleJobError::from_anyhow_error(anyhow!("could not download VAE: {:?}", err))
})?;
let vae_weight_record = get_weight_by_token(
&model_weight_vae,
false,
&deps.db.mysql_pool
).await?;
let vae_weight_record = match vae_weight_record {
Some(val) => val,
None => {
return Err(
ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))
);
}
};
let vae_details = RemoteCloudBucketDetails::new(
vae_weight_record.public_bucket_hash.clone(),
vae_weight_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()),
vae_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string())
);
remote_cloud_file_client.download_file(vae_details, path_to_string(vae_path.clone())).await?;
// // VAE token for now
// let vae_token = String::from("REPLACE_ME");
// let model_weight_vae = ModelWeightToken(vae_token);
//
// let vae_weight_record = get_weight_by_token(
// &model_weight_vae,
// false,
// &deps.db.mysql_pool
// ).await?;
//
// let vae_weight_record = match vae_weight_record {
// Some(val) => val,
// None => {
// return Err(
// ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))
// );
// }
// };
//
// let vae_details = RemoteCloudBucketDetails::new(
// vae_weight_record.public_bucket_hash.clone(),
// vae_weight_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()),
// vae_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string())
// );
//
// remote_cloud_file_client.download_file(vae_details, path_to_string(vae_path.clone())).await?;
let prompt = match sd_args.maybe_prompt {
Some(val) => val,
crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs
+2
View File
@@ -7,6 +7,7 @@ use crate::job::job_types::image_generation::sd::sd_inference_command::StableDif
pub struct StableDiffusionDependencies {
pub downloaders: StableDiffusionDownloaders,
pub inference_command: StableDiffusionInferenceCommand,
pub vae_bucket_path: String,
}
impl StableDiffusionDependencies {
@@ -14,6 +15,7 @@ impl StableDiffusionDependencies {
Ok(Self {
downloaders: StableDiffusionDownloaders::build_all_from_env(),
inference_command: StableDiffusionInferenceCommand::from_env()?,
vae_bucket_path: easyenv::get_env_string_required("SD_VAE_BUCKET_PATH")?,
})
}
}