diff --git a/crates/schema/public/enums/src/common/model_type.rs b/crates/schema/public/enums/src/common/model_type.rs index 3569b49100..6210791959 100644 --- a/crates/schema/public/enums/src/common/model_type.rs +++ b/crates/schema/public/enums/src/common/model_type.rs @@ -12,9 +12,37 @@ use utoipa::ToSchema; #[derive(Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd, Serialize, Deserialize, ToSchema)] #[serde(rename_all = "snake_case")] pub enum ModelType { - Artcraft, - Fal, - Sora, + // Image models + #[serde(rename = "flux_1_dev")] + Flux1Dev, + #[serde(rename = "flux_1_schnell")] + Flux1Schnell, + #[serde(rename = "flux_pro_1p1")] + FluxPro11, + #[serde(rename = "flux_pro_1p1_ultra")] + FluxPro11Ultra, + #[serde(rename = "gpt_image_1")] + GptImage1, + #[serde(rename = "recraft_3")] + Recraft3, + + // Video models + #[serde(rename = "kling_1p6_pro")] + Kling16Pro, + #[serde(rename = "kling_2p1_pro")] + Kling21Pro, + #[serde(rename = "kling_2p1_master")] + Kling21Master, + #[serde(rename = "seedance_1p0_lite")] + Seedance10Lite, + #[serde(rename = "veo_2")] + Veo2, + + // 3D Object generation models + #[serde(rename = "hunyuan_3d_2p0")] + Hunyuan3d2_0, + #[serde(rename = "hunyuan_3d_2p1")] + Hunyuan3d2_1, } impl_enum_display_and_debug_using_to_str!(ModelType); @@ -26,17 +54,48 @@ impl_mysql_from_row!(ModelType); impl ModelType { pub fn to_str(&self) -> &'static str { match self { - Self::Artcraft => "artcraft", - Self::Fal => "fal", - Self::Sora => "sora", + // Image models + Self::Flux1Dev => "flux_1_dev", + Self::Flux1Schnell => "flux_1_schnell", + Self::FluxPro11 => "flux_pro_1p1", + Self::FluxPro11Ultra => "flux_pro_1p1_ultra", + Self::GptImage1 => "gpt_image_1", + Self::Recraft3 => "recraft_3", + + // Video models + Self::Kling16Pro => "kling_1p6_pro", + Self::Kling21Pro => "kling_2p1_pro", + Self::Kling21Master => "kling_2p1_master", + Self::Seedance10Lite => "seedance_1p0_lite", + Self::Veo2 => "veo_2", + + // 3D Object generation models + Self::Hunyuan3d2_0 => "hunyuan_3d_2p0", + Self::Hunyuan3d2_1 => "hunyuan_3d_2p1", } } pub fn from_str(job_status: &str) -> Result { match job_status { - "artcraft" => Ok(Self::Artcraft), - "fal" => Ok(Self::Fal), - "sora" => Ok(Self::Sora), + // Image models + "flux_1_dev" => Ok(Self::Flux1Dev), + "flux_1_schnell" => Ok(Self::Flux1Schnell), + "flux_pro_1p1" => Ok(Self::FluxPro11), + "flux_pro_1p1_ultra" => Ok(Self::FluxPro11Ultra), + "gpt_image_1" => Ok(Self::GptImage1), + "recraft_3" => Ok(Self::Recraft3), + + // Video models + "kling_1p6_pro" => Ok(Self::Kling16Pro), + "kling_2p1_pro" => Ok(Self::Kling21Pro), + "kling_2p1_master" => Ok(Self::Kling21Master), + "seedance_1p0_lite" => Ok(Self::Seedance10Lite), + "veo_2" => Ok(Self::Veo2), + + // 3D Object generation models + "hunyuan_3d_2p0" => Ok(Self::Hunyuan3d2_0), + "hunyuan_3d_2p1" => Ok(Self::Hunyuan3d2_1), + _ => Err(format!("invalid task_state: {:?}", job_status)), } } @@ -45,9 +104,24 @@ impl ModelType { // NB: BTreeSet is sorted // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. BTreeSet::from([ - Self::Artcraft, - Self::Fal, - Self::Sora, + // Image models + Self::Flux1Dev, + Self::Flux1Schnell, + Self::FluxPro11, + Self::FluxPro11Ultra, + Self::GptImage1, + Self::Recraft3, + + // Video models + Self::Kling16Pro, + Self::Kling21Pro, + Self::Kling21Master, + Self::Seedance10Lite, + Self::Veo2, + + // 3D Object generation models + Self::Hunyuan3d2_0, + Self::Hunyuan3d2_1, ]) } } @@ -62,32 +136,87 @@ mod tests { #[test] fn test_serialization() { - assert_serialization(ModelType::Artcraft, "artcraft"); - assert_serialization(ModelType::Fal, "fal"); - assert_serialization(ModelType::Sora, "sora"); + // Image models + assert_serialization(ModelType::Flux1Dev, "flux_1_dev"); + assert_serialization(ModelType::Flux1Schnell, "flux_1_schnell"); + assert_serialization(ModelType::FluxPro11, "flux_pro_1p1"); + assert_serialization(ModelType::FluxPro11Ultra, "flux_pro_1p1_ultra"); + assert_serialization(ModelType::GptImage1, "gpt_image_1"); + assert_serialization(ModelType::Recraft3, "recraft_3"); + // Video models + assert_serialization(ModelType::Kling16Pro, "kling_1p6_pro"); + assert_serialization(ModelType::Kling21Pro, "kling_2p1_pro"); + assert_serialization(ModelType::Kling21Master, "kling_2p1_master"); + assert_serialization(ModelType::Seedance10Lite, "seedance_1p0_lite"); + assert_serialization(ModelType::Veo2, "veo_2"); + // 3D Object generation models + assert_serialization(ModelType::Hunyuan3d2_0, "hunyuan_3d_2p0"); + assert_serialization(ModelType::Hunyuan3d2_1, "hunyuan_3d_2p1"); } #[test] fn to_str() { - assert_eq!(ModelType::Artcraft.to_str(), "artcraft"); - assert_eq!(ModelType::Fal.to_str(), "fal"); - assert_eq!(ModelType::Sora.to_str(), "sora"); + // Image models + assert_eq!(ModelType::Flux1Dev.to_str(), "flux_1_dev"); + assert_eq!(ModelType::Flux1Schnell.to_str(), "flux_1_schnell"); + assert_eq!(ModelType::FluxPro11.to_str(), "flux_pro_1p1"); + assert_eq!(ModelType::FluxPro11Ultra.to_str(), "flux_pro_1p1_ultra"); + assert_eq!(ModelType::GptImage1.to_str(), "gpt_image_1"); + assert_eq!(ModelType::Recraft3.to_str(), "recraft_3"); + + // Video models + assert_eq!(ModelType::Kling16Pro.to_str(), "kling_1p6_pro"); + assert_eq!(ModelType::Kling21Pro.to_str(), "kling_2p1_pro"); + assert_eq!(ModelType::Kling21Master.to_str(), "kling_2p1_master"); + assert_eq!(ModelType::Seedance10Lite.to_str(), "seedance_1p0_lite"); + assert_eq!(ModelType::Veo2.to_str(), "veo_2"); + + // 3D Object generation models + assert_eq!(ModelType::Hunyuan3d2_0.to_str(), "hunyuan_3d_2p0"); + assert_eq!(ModelType::Hunyuan3d2_1.to_str(), "hunyuan_3d_2p1"); } #[test] fn from_str() { - assert_eq!(ModelType::from_str("artcraft").unwrap(), ModelType::Artcraft); - assert_eq!(ModelType::from_str("fal").unwrap(), ModelType::Fal); - assert_eq!(ModelType::from_str("sora").unwrap(), ModelType::Sora); + // Image models + assert_eq!(ModelType::from_str("flux_1_dev").unwrap(), ModelType::Flux1Dev); + assert_eq!(ModelType::from_str("flux_1_schnell").unwrap(), ModelType::Flux1Schnell); + assert_eq!(ModelType::from_str("flux_pro_1p1").unwrap(), ModelType::FluxPro11); + assert_eq!(ModelType::from_str("flux_pro_1p1_ultra").unwrap(), ModelType::FluxPro11Ultra); + assert_eq!(ModelType::from_str("gpt_image_1").unwrap(), ModelType::GptImage1); + assert_eq!(ModelType::from_str("recraft_3").unwrap(), ModelType::Recraft3); + // Video models + assert_eq!(ModelType::from_str("kling_1p6_pro").unwrap(), ModelType::Kling16Pro); + assert_eq!(ModelType::from_str("kling_2p1_pro").unwrap(), ModelType::Kling21Pro); + assert_eq!(ModelType::from_str("kling_2p1_master").unwrap(), ModelType::Kling21Master); + assert_eq!(ModelType::from_str("seedance_1p0_lite").unwrap(), ModelType::Seedance10Lite); + assert_eq!(ModelType::from_str("veo_2").unwrap(), ModelType::Veo2); + // 3D Object generation models + assert_eq!(ModelType::from_str("hunyuan_3d_2p0").unwrap(), ModelType::Hunyuan3d2_0); + assert_eq!(ModelType::from_str("hunyuan_3d_2p1").unwrap(), ModelType::Hunyuan3d2_1); } #[test] fn all_variants() { let mut variants = ModelType::all_variants(); - assert_eq!(variants.len(), 3); - assert_eq!(variants.pop_first(), Some(ModelType::Artcraft)); - assert_eq!(variants.pop_first(), Some(ModelType::Fal)); - assert_eq!(variants.pop_first(), Some(ModelType::Sora)); + assert_eq!(variants.len(), 13); + // Image models + assert_eq!(variants.pop_first(), Some(ModelType::Flux1Dev)); + assert_eq!(variants.pop_first(), Some(ModelType::Flux1Schnell)); + assert_eq!(variants.pop_first(), Some(ModelType::FluxPro11)); + assert_eq!(variants.pop_first(), Some(ModelType::FluxPro11Ultra)); + assert_eq!(variants.pop_first(), Some(ModelType::GptImage1)); + assert_eq!(variants.pop_first(), Some(ModelType::Recraft3)); + // Video models + assert_eq!(variants.pop_first(), Some(ModelType::Kling16Pro)); + assert_eq!(variants.pop_first(), Some(ModelType::Kling21Pro)); + assert_eq!(variants.pop_first(), Some(ModelType::Kling21Master)); + assert_eq!(variants.pop_first(), Some(ModelType::Seedance10Lite)); + assert_eq!(variants.pop_first(), Some(ModelType::Veo2)); + // 3D Object generation models + assert_eq!(variants.pop_first(), Some(ModelType::Hunyuan3d2_0)); + assert_eq!(variants.pop_first(), Some(ModelType::Hunyuan3d2_1)); + assert_eq!(variants.pop_first(), None); } }