mirror of
https://github.com/storytold/artcraft.git
synced 2026-11-04 09:28:11 +00:00
model type enum
This commit is contained in:
Brandon Thomas committed
1 parent
77fbd84006
commit
fb68554f85
1 file changed
+154
-25
@@ -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<Self, String> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user