From 2b100f2aaad7a7244d8f746ffbb30ff578c82376 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 16 Aug 2024 20:40:56 -0400 Subject: [PATCH] replace old SD json metadata with new polymorphic enum struct --- .../inner_payloads/mod.rs | 3 +- .../stable_diffusion_extra_info.rs | 72 +++++++++++++++++++ .../media_file_extra_info.rs | 22 +++++- .../insert_media_file_from_gptsovits.rs | 2 +- .../insert_media_file_from_live_portrait.rs | 2 +- .../create/insert_media_file_generic.rs | 18 +++-- .../fbx_to_gltf/process_job.rs | 2 +- .../image_generation/sd/process_job.rs | 14 ---- .../sd/process_job_inference.rs | 33 ++++----- .../process_job.rs | 2 +- 10 files changed, 126 insertions(+), 44 deletions(-) create mode 100644 crates/lib/mysql_queries/src/payloads/media_file_extra_info/inner_payloads/stable_diffusion_extra_info.rs diff --git a/crates/lib/mysql_queries/src/payloads/media_file_extra_info/inner_payloads/mod.rs b/crates/lib/mysql_queries/src/payloads/media_file_extra_info/inner_payloads/mod.rs index f33eeeaa70..6a0b5f1765 100644 --- a/crates/lib/mysql_queries/src/payloads/media_file_extra_info/inner_payloads/mod.rs +++ b/crates/lib/mysql_queries/src/payloads/media_file_extra_info/inner_payloads/mod.rs @@ -1 +1,2 @@ -pub mod live_portrait_video_extra_info; \ No newline at end of file +pub mod live_portrait_video_extra_info; +pub mod stable_diffusion_extra_info; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/payloads/media_file_extra_info/inner_payloads/stable_diffusion_extra_info.rs b/crates/lib/mysql_queries/src/payloads/media_file_extra_info/inner_payloads/stable_diffusion_extra_info.rs new file mode 100644 index 0000000000..4c2b2620e7 --- /dev/null +++ b/crates/lib/mysql_queries/src/payloads/media_file_extra_info/inner_payloads/stable_diffusion_extra_info.rs @@ -0,0 +1,72 @@ + +/// We used to write stable diffusion outputs with these values in +/// the `extra_file_modification_info` field. +#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct StableDiffusionExtraInfo { + // NB: If we decide to encode this information again in the future, + // it's probably best to use different fields or an entirely different struct. + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub cfg_scale: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub negative_prompt: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub lora_model_weight_token: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub lora_name: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub sampler: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub width: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub height: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub seed: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub number_of_samples: Option, +} + +#[cfg(test)] +mod tests { + use crate::payloads::media_file_extra_info::inner_payloads::stable_diffusion_extra_info::StableDiffusionExtraInfo; + use crate::payloads::media_file_extra_info::media_file_extra_info::MediaFileExtraInfo; + + #[test] + fn wrapped_enum() { + let payload = MediaFileExtraInfo::S(StableDiffusionExtraInfo { + prompt: Some("foo bar baz".to_string()), + ..Default::default() + }); + + let json = r#"{"S":{"prompt":"foo bar baz"}}"#; + + let serialized = payload.to_string().unwrap(); + assert_eq!(&serialized, json); + + let deserialized = MediaFileExtraInfo::from_str(&json).unwrap(); + assert_eq!(deserialized, payload); + } + + #[test] + fn unwrapped_enum() { + let json = r#"{"prompt":"foo bar baz"}"#; + let deserialized = MediaFileExtraInfo::from_str(&json).unwrap(); + + let expected = MediaFileExtraInfo::S(StableDiffusionExtraInfo { + prompt: Some("foo bar baz".to_string()), + ..Default::default() + }); + + assert_eq!(deserialized, expected); + } +} diff --git a/crates/lib/mysql_queries/src/payloads/media_file_extra_info/media_file_extra_info.rs b/crates/lib/mysql_queries/src/payloads/media_file_extra_info/media_file_extra_info.rs index b7eeefc94e..a610b8438f 100644 --- a/crates/lib/mysql_queries/src/payloads/media_file_extra_info/media_file_extra_info.rs +++ b/crates/lib/mysql_queries/src/payloads/media_file_extra_info/media_file_extra_info.rs @@ -1,5 +1,8 @@ +use anyhow::{anyhow, bail}; use errors::AnyhowResult; + use crate::payloads::media_file_extra_info::inner_payloads::live_portrait_video_extra_info::LivePortraitVideoExtraInfo; +use crate::payloads::media_file_extra_info::inner_payloads::stable_diffusion_extra_info::StableDiffusionExtraInfo; /// For things that don't fit in the schema #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] @@ -8,11 +11,28 @@ pub enum MediaFileExtraInfo { /// NB: Enum variant is short to conserve DB space. /// NB: DO NOT CHANGE. It could break live jobs. L(LivePortraitVideoExtraInfo), + + /// Legacy Stable Diffusion info we used to write (though we + /// did not use nested polymorphic enums structs for the JSON) + /// If we do stable diffusion again, it's probably best to + /// design a new and more efficient/compact struct. + S(StableDiffusionExtraInfo), } impl MediaFileExtraInfo { pub fn from_str(value: &str) -> AnyhowResult { - Ok(serde_json::from_str(value)?) + let result = serde_json::from_str(value); + let err = match result { + Ok(value) => return Ok(value), + Err(err) => err, + }; + // NB: Some older stable diffusion payloads were raw equivalents of `StableDiffusionExtraInfo`. + if value.contains("prompt") { + let stable_diffusion_info : StableDiffusionExtraInfo = serde_json::from_str(value)?; + return Ok(Self::S(stable_diffusion_info)); + } else { + return bail!(err); + } } pub fn to_string(&self) -> AnyhowResult { diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_gptsovits.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_gptsovits.rs index d5e1d64716..da26228ffe 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_gptsovits.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_gptsovits.rs @@ -86,7 +86,7 @@ pub async fn insert_media_file_from_gptsovits( maybe_frame_width: None, maybe_frame_height: None, maybe_prompt_token: None, - extra_file_modification_info: None, + maybe_extra_media_info: None, maybe_mod_user_token: None, }).await?; diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_live_portrait.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_live_portrait.rs index 0b4a5b2ae2..414d36f817 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_live_portrait.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_live_portrait.rs @@ -89,7 +89,7 @@ pub async fn insert_media_file_from_live_portrait( maybe_origin_filename: None, maybe_batch_token: None, maybe_prompt_token: None, - extra_file_modification_info: None, // TODO ??? + maybe_extra_media_info: None, // TODO maybe_mod_user_token: None, }).await?; diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs index 87e6191d81..e61cf31827 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs @@ -14,7 +14,7 @@ use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::prompts::PromptToken; use tokens::tokens::users::UserToken; - +use crate::payloads::media_file_extra_info::media_file_extra_info::MediaFileExtraInfo; use crate::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; use crate::queries::generic_synthetic_ids::transactional_increment_generic_synthetic_id::transactional_increment_generic_synthetic_id; @@ -54,11 +54,15 @@ pub struct InsertArgs<'a> { pub public_bucket_directory_hash: &'a str, pub maybe_public_bucket_prefix: Option<&'a str>, pub maybe_public_bucket_extension: Option<&'a str>, - pub extra_file_modification_info: Option<&'a str>, // Assuming TEXT type can be represented by &str pub maybe_creator_file_synthetic_id_category: IdCategory, pub maybe_creator_category_synthetic_id_category: IdCategory, + /// Extra polymorphic data stored in `extra_file_modification_info` + /// This differs on a per media type basis and can depend on the product + /// that generates the media file. + pub maybe_extra_media_info: Option<&'a MediaFileExtraInfo>, + // Modification and generation info pub maybe_mod_user_token: Option<&'a UserToken>, pub is_generated_on_prem: bool, @@ -72,6 +76,10 @@ pub async fn insert_media_file_generic( { let result_token = MediaFileToken::generate(); + let extra_file_modification_info = args + .maybe_extra_media_info.map(|extra| extra.to_string()) + .transpose()?; + let mut maybe_creator_file_synthetic_id : Option = None; let mut maybe_creator_category_synthetic_id : Option = None; @@ -133,7 +141,6 @@ pub async fn insert_media_file_generic( public_bucket_directory_hash = ?, maybe_public_bucket_prefix = ?, maybe_public_bucket_extension = ?, - extra_file_modification_info = ?, maybe_creator_user_token = ?, maybe_creator_anonymous_visitor_token = ?, @@ -144,6 +151,8 @@ pub async fn insert_media_file_generic( maybe_creator_file_synthetic_id = ?, maybe_creator_category_synthetic_id = ?, + extra_file_modification_info = ?, + maybe_mod_user_token = ?, is_generated_on_prem = ?, generated_by_worker = ?, @@ -179,7 +188,6 @@ pub async fn insert_media_file_generic( args.public_bucket_directory_hash, args.maybe_public_bucket_prefix, args.maybe_public_bucket_extension, - args.extra_file_modification_info, args.job.maybe_creator_user_token.as_deref(), args.job.maybe_creator_anonymous_visitor_token.as_deref(), @@ -190,6 +198,8 @@ pub async fn insert_media_file_generic( maybe_creator_file_synthetic_id, maybe_creator_category_synthetic_id, + extra_file_modification_info, + args.maybe_mod_user_token, args.is_generated_on_prem, args.generated_by_worker, diff --git a/crates/service/job/inference_job/src/job/job_types/format_conversion/fbx_to_gltf/process_job.rs b/crates/service/job/inference_job/src/job/job_types/format_conversion/fbx_to_gltf/process_job.rs index a5fdbe6d79..d6e779e654 100644 --- a/crates/service/job/inference_job/src/job/job_types/format_conversion/fbx_to_gltf/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/format_conversion/fbx_to_gltf/process_job.rs @@ -204,7 +204,7 @@ pub async fn process_job(args: FbxToGltfJobArgs<'_>) -> Result, - pub lora_model_weight_token: Option, - pub lora_name: Option, - pub sampler: String, - pub width: u32, - pub height: u32, - pub seed: i64, - pub number_of_samples: u32, -} - pub struct StableDiffusionProcessArgs<'a> { pub job_dependencies: &'a JobDependencies, pub job: &'a AvailableInferenceJob, diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job_inference.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job_inference.rs index e5c6479a26..91acf55367 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job_inference.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job_inference.rs @@ -17,6 +17,8 @@ use enums::by_table::media_files::media_file_origin_product_category::MediaFileO use enums::by_table::media_files::media_file_type::MediaFileType; use enums::by_table::prompts::prompt_type::PromptType; use filesys::path_to_string::path_to_string; +use mysql_queries::payloads::media_file_extra_info::inner_payloads::stable_diffusion_extra_info::StableDiffusionExtraInfo; +use mysql_queries::payloads::media_file_extra_info::media_file_extra_info::MediaFileExtraInfo; use mysql_queries::queries::media_files::create::insert_media_file_generic::{insert_media_file_generic, InsertArgs}; use mysql_queries::queries::model_weights::get::get_weight::get_weight_by_token; use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs}; @@ -25,7 +27,7 @@ use tokens::tokens::prompts::PromptToken; use crate::job::job_loop::job_success_result::{JobSuccessResult, ResultEntity}; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; -use crate::job::job_types::image_generation::sd::process_job::{InferenceValues, sd_args_from_job, StableDiffusionProcessArgs}; +use crate::job::job_types::image_generation::sd::process_job::{sd_args_from_job, StableDiffusionProcessArgs}; use crate::job::job_types::image_generation::sd::sd_inference_command::InferenceArgs; pub async fn process_job_inference( @@ -212,27 +214,18 @@ pub async fn process_job_inference( let mut entries = vec![]; - let inputs = InferenceValues { - prompt: prompt.clone(), - cfg_scale: sd_args.maybe_cfg_scale.unwrap_or(7), + let inputs = MediaFileExtraInfo::S(StableDiffusionExtraInfo { + prompt: Some(prompt.clone()), + cfg_scale: Some(sd_args.maybe_cfg_scale.unwrap_or(7)), negative_prompt: sd_args.maybe_n_prompt, lora_model_weight_token: Some(lora_token), lora_name: Some(lora_name), - sampler: sd_args.maybe_sampler.unwrap_or(String::from("Euler a")), - width: sd_args.maybe_width.unwrap_or(512), - height: sd_args.maybe_height.unwrap_or(512), - seed: sd_args.maybe_seed.unwrap_or(1), - number_of_samples, - }; - - let inputs = match serde_json::to_string(&inputs) { - Ok(result) => result, - Err(_err) => { - return Err( - ProcessSingleJobError::from_anyhow_error(anyhow!("couldn't serialize metadata.")) - ); - } - }; + sampler: Some(sd_args.maybe_sampler.unwrap_or(String::from("Euler a"))), + width: Some(sd_args.maybe_width.unwrap_or(512)), + height: Some(sd_args.maybe_height.unwrap_or(512)), + seed: Some(sd_args.maybe_seed.unwrap_or(1)), + number_of_samples: Some(number_of_samples), + }); let batch_token = BatchGenerationToken::generate(); let prompt_token = PromptToken::generate(); @@ -285,7 +278,7 @@ pub async fn process_job_inference( public_bucket_directory_hash: bucket_details.object_hash.as_str(), maybe_public_bucket_prefix: Some(bucket_details.prefix.as_str()), maybe_public_bucket_extension: Some(bucket_details.suffix.as_str()), - extra_file_modification_info: Some(&inputs), + maybe_extra_media_info: Some(&inputs), maybe_creator_file_synthetic_id_category: IdCategory::MediaFile, maybe_creator_category_synthetic_id_category: IdCategory::ModelWeights, maybe_mod_user_token: None, diff --git a/crates/service/job/inference_job/src/job/job_types/render_engine_scene/render_engine_scene_to_video/process_job.rs b/crates/service/job/inference_job/src/job/job_types/render_engine_scene/render_engine_scene_to_video/process_job.rs index 0645f74bde..d68af7a008 100644 --- a/crates/service/job/inference_job/src/job/job_types/render_engine_scene/render_engine_scene_to_video/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/render_engine_scene/render_engine_scene_to_video/process_job.rs @@ -273,7 +273,7 @@ pub async fn process_job(args: BvhToWorkflowJobArgs<'_>) -> Result