diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs index 19bf5a74c7..99c146683f 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs @@ -28,6 +28,9 @@ pub enum InferenceCategory { #[serde(rename = "video_filter")] VideoFilter, + + #[serde(rename = "image_generation")] + ImageGeneration } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -42,6 +45,7 @@ impl InferenceCategory { Self::TextToSpeech => "text_to_speech", Self::VoiceConversion => "voice_conversion", Self::VideoFilter => "video_filter", + Self::ImageGeneration => "image_generation", } } @@ -51,6 +55,7 @@ impl InferenceCategory { "text_to_speech" => Ok(Self::TextToSpeech), "voice_conversion" => Ok(Self::VoiceConversion), "video_filter" => Ok(Self::VideoFilter), + "image_generation" => Ok(Self::ImageGeneration), _ => Err(format!("invalid value: {:?}", value)), } } @@ -63,6 +68,7 @@ impl InferenceCategory { Self::TextToSpeech, Self::VoiceConversion, Self::VideoFilter, + Self::ImageGeneration, ]) } } @@ -78,6 +84,7 @@ mod tests { assert_serialization(InferenceCategory::TextToSpeech, "text_to_speech"); assert_serialization(InferenceCategory::VoiceConversion, "voice_conversion"); assert_serialization(InferenceCategory::VideoFilter, "video_filter"); + assert_serialization(InferenceCategory::ImageGeneration, "image_generation"); } #[test] @@ -86,6 +93,7 @@ mod tests { assert_eq!(InferenceCategory::TextToSpeech.to_str(), "text_to_speech"); assert_eq!(InferenceCategory::VoiceConversion.to_str(), "voice_conversion"); assert_eq!(InferenceCategory::VideoFilter.to_str(), "video_filter"); + assert_eq!(InferenceCategory::ImageGeneration.to_str(), "image_generation"); } #[test] @@ -94,6 +102,7 @@ mod tests { assert_eq!(InferenceCategory::from_str("text_to_speech").unwrap(), InferenceCategory::TextToSpeech); assert_eq!(InferenceCategory::from_str("voice_conversion").unwrap(), InferenceCategory::VoiceConversion); assert_eq!(InferenceCategory::from_str("video_filter").unwrap(), InferenceCategory::VideoFilter); + assert_eq!(InferenceCategory::from_str("image_generation").unwrap(), InferenceCategory::ImageGeneration); } #[test] @@ -105,6 +114,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(InferenceCategory::TextToSpeech)); assert_eq!(variants.pop_first(), Some(InferenceCategory::VoiceConversion)); assert_eq!(variants.pop_first(), Some(InferenceCategory::VideoFilter)); + assert_eq!(variants.pop_first(), Some(InferenceCategory::ImageGeneration)); assert_eq!(variants.pop_first(), None); // Generated check diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs index 49da8fb98c..90bc808db2 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs @@ -3,10 +3,10 @@ use errors::AnyhowResult; use crate::payloads::generic_inference_args::lipsync_payload::LipsyncArgs; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; -use crate::payloads::generic_inference_args::videofilter_payload::{RerenderArgs}; -use crate::payloads::generic_inference_args::image_generation_payload::{SDArgs}; - +use crate::payloads::generic_inference_args::videofilter_payload::RerenderArgs; +use crate::payloads::generic_inference_args::image_generation_payload::SDArgs; +use tokens::tokens::{model_weights::ModelWeightToken, media_files::MediaFileToken}; /// Used to encode extra state for the `generic_inference_jobs` table. /// This should act somewhat like a serialized protobuf stored inside a record. #[derive(Clone, Debug, Serialize, Deserialize)] @@ -96,7 +96,10 @@ pub enum PolymorphicInferenceArgs { }, /// Rerender a video. (Short name to save space when serializing.) - Rr(RerenderArgs) + Rr(RerenderArgs), + + /// Image generation. (Short name to save space when serializing.) + Ig(SDArgs) } @@ -135,9 +138,13 @@ impl InferenceCategoryAbbreviated { #[cfg(test)] mod tests { - use crate::payloads::generic_inference_args::generic_inference_args::{FundamentalFrequencyMethodForJob, GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; + use enums::by_table::generic_inference_jobs::inference_category; + +use crate::payloads::generic_inference_args::generic_inference_args::{FundamentalFrequencyMethodForJob, GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; use crate::payloads::generic_inference_args::lipsync_payload::{LipsyncAnimationAudioSource, LipsyncAnimationImageSource, LipsyncArgs}; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; + use crate::payloads::generic_inference_args::image_generation_payload::SDArgs; + use tokens::tokens::{model_weights::ModelWeightToken, media_files::MediaFileToken}; #[test] fn typical_lipsync_animation_args_serialize() { @@ -184,6 +191,44 @@ mod tests { assert!(json.len() < 1000); } + #[test] + fn typical_image_gen_args_serialize() { + + let video_media_token = MediaFileToken("video_media_token".to_string()); + let image_media_token = MediaFileToken("image_media_token".to_string()); + let sd_model_token = ModelWeightToken("sd_model_token".to_string()); + let lora_model_token = ModelWeightToken("lora_model_token".to_string()); + + let prompt = "prompt".to_string(); + let a_prompt = "a_prompt".to_string(); + let n_prompt = "n_prompt".to_string(); + let seed = 1; + let upload_path = "upload_path".to_string(); + let lora_upload_path = "lora_upload_path".to_string(); + + let args = GenericInferenceArgs { + inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration), + args: Some(PolymorphicInferenceArgs::Ig(SDArgs { + maybe_video_source: Some(video_media_token), + maybe_image_source: Some(image_media_token), + maybe_sd_model_token: Some(sd_model_token), + maybe_lora_model_token: Some(lora_model_token), + maybe_prompt: Some(prompt), + maybe_a_prompt: Some(a_prompt), + maybe_n_prompt: Some(n_prompt), + maybe_seed: Some(seed), + maybe_upload_path: Some(upload_path), + maybe_lora_upload_path: Some(lora_upload_path) + })), + }; + + let json = serde_json::ser::to_string(&args).unwrap(); + + assert_eq!(json, r#"{"cat":"ig","args":{"Ig":{"v":"video_media_token","i":"image_media_token","s":"sd_model_token","l":"lora_model_token","p":"prompt","a":"a_prompt","n":"n_prompt","e":1,"u":"upload_path","o":"lora_upload_path"}}}"#.to_string()); + + + } + #[test] fn typical_voice_conversion_args_serialize() { let args = GenericInferenceArgs { diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs index 47c1efa11b..87a23e9969 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs @@ -1,4 +1,4 @@ -use tokens::tokens::model_weights::ModelWeightToken; +use tokens::tokens::{model_weights::ModelWeightToken, media_files::MediaFileToken}; /// Video sources can be one of several: /// - F: media_files (todo) @@ -22,7 +22,11 @@ pub enum SDSource { pub struct SDArgs { #[serde(rename = "vs")] #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_video_source: Option, + pub maybe_video_source: Option, + + #[serde(rename = "is")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_image_source: Option, #[serde(rename = "sd")] #[serde(skip_serializing_if = "Option::is_none")] @@ -47,8 +51,14 @@ pub struct SDArgs { #[serde(rename = "se")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_seed: Option, - - #[serde(rename = "upload_path")] + + #[serde(rename = "mup")] #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_upload_path: Option + pub maybe_upload_path: Option, + + #[serde(rename = "lup")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_lora_upload_path: Option } + +