From 9326a856cd066325787bf9209cbae0899cdb6ce6 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 15 Aug 2025 23:22:12 -0400 Subject: [PATCH] recognize more midjourney models --- .../src/endpoints/imagine.rs | 6 +++ .../threads/midjourney_long_polling_thread.rs | 3 ++ .../public/enums/src/common/model_type.rs | 47 +++++++++++++++---- .../libs/model-list/src/lib/ModelMapping.ts | 17 +++++++ 4 files changed, 63 insertions(+), 10 deletions(-) diff --git a/crates/api_clients/midjourney_client/src/endpoints/imagine.rs b/crates/api_clients/midjourney_client/src/endpoints/imagine.rs index e505f7e0eb..35ea862e20 100644 --- a/crates/api_clients/midjourney_client/src/endpoints/imagine.rs +++ b/crates/api_clients/midjourney_client/src/endpoints/imagine.rs @@ -39,7 +39,10 @@ pub struct ImagineItem { #[derive(Debug, Clone)] pub enum MidjourneyJobType { V6Diffusion, + V6p1Diffusion, + V6p1RawDiffusion, V7Diffusion, + V7DraftDiffusion, V7DraftRawDiffusion, V7RawDiffusion, Vid11I2vRenderAJointVideo, @@ -53,7 +56,10 @@ impl MidjourneyJobType { pub fn from_str(s: &str) -> MidjourneyJobType { match s { "v6_diffusion" => MidjourneyJobType::V6Diffusion, + "v6-1_diffusion" => MidjourneyJobType::V6p1Diffusion, + "v6-1_raw_diffusion" => MidjourneyJobType::V6p1RawDiffusion, "v7_diffusion" => MidjourneyJobType::V7Diffusion, + "v7_draft_diffusion" => MidjourneyJobType::V7DraftDiffusion, "v7_draft_raw_diffusion" => MidjourneyJobType::V7DraftRawDiffusion, "v7_raw_diffusion" => MidjourneyJobType::V7RawDiffusion, "vid_1.1_i2v_render_a_joint_video" => MidjourneyJobType::Vid11I2vRenderAJointVideo, diff --git a/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs b/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs index 862c6a5eed..99ed48cdbb 100644 --- a/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs +++ b/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs @@ -203,8 +203,11 @@ async fn check_midjourney_tasks( let model_type = match midjourney_item.job_type { Some(MidjourneyJobType::V6Diffusion) => ModelType::MidjourneyV6, + Some(MidjourneyJobType::V6p1Diffusion) => ModelType::MidjourneyV6p1, + Some(MidjourneyJobType::V6p1RawDiffusion) => ModelType::MidjourneyV6p1Raw, Some(MidjourneyJobType::V7Diffusion) => ModelType::MidjourneyV7, Some(MidjourneyJobType::V7RawDiffusion) => ModelType::MidjourneyV7Raw, + Some(MidjourneyJobType::V7DraftDiffusion) => ModelType::MidjourneyV7Draft, Some(MidjourneyJobType::V7DraftRawDiffusion) => ModelType::MidjourneyV7DraftRaw, Some(MidjourneyJobType::Other(ref other)) => { info!("Unknown Midjourney job type (for job id {}): {}", midjourney_job_id, other); diff --git a/crates/schema/public/enums/src/common/model_type.rs b/crates/schema/public/enums/src/common/model_type.rs index b3af134eff..860c99ce0a 100644 --- a/crates/schema/public/enums/src/common/model_type.rs +++ b/crates/schema/public/enums/src/common/model_type.rs @@ -35,12 +35,18 @@ pub enum ModelType { Midjourney, #[serde(rename = "midjourney_v6")] MidjourneyV6, + #[serde(rename = "midjourney_v6p1")] + MidjourneyV6p1, + #[serde(rename = "midjourney_v6p1_raw")] + MidjourneyV6p1Raw, #[serde(rename = "midjourney_v7")] MidjourneyV7, - #[serde(rename = "midjourney_v7_raw")] - MidjourneyV7Raw, + #[serde(rename = "midjourney_v7_draft")] + MidjourneyV7Draft, #[serde(rename = "midjourney_v7_draft_raw")] MidjourneyV7DraftRaw, + #[serde(rename = "midjourney_v7_raw")] + MidjourneyV7Raw, //// Image Infill models //#[serde(rename = "flux_pro_1_infill")] @@ -85,9 +91,12 @@ impl ModelType { Self::Recraft3 => "recraft_3", Self::Midjourney => "midjourney", Self::MidjourneyV6 => "midjourney_v6", + Self::MidjourneyV6p1 => "midjourney_v6p1", + Self::MidjourneyV6p1Raw => "midjourney_v6p1_raw", Self::MidjourneyV7 => "midjourney_v7", - Self::MidjourneyV7Raw => "midjourney_v7_raw", + Self::MidjourneyV7Draft => "midjourney_v7_draft", Self::MidjourneyV7DraftRaw => "midjourney_v7_draft_raw", + Self::MidjourneyV7Raw => "midjourney_v7_raw", // Video models Self::Kling16Pro => "kling_1p6_pro", @@ -115,9 +124,12 @@ impl ModelType { "recraft_3" => Ok(Self::Recraft3), "midjourney" => Ok(Self::Midjourney), "midjourney_v6" => Ok(Self::MidjourneyV6), + "midjourney_v6p1" => Ok(Self::MidjourneyV6p1), + "midjourney_v6p1_raw" => Ok(Self::MidjourneyV6p1Raw), "midjourney_v7" => Ok(Self::MidjourneyV7), - "midjourney_v7_raw" => Ok(Self::MidjourneyV7Raw), + "midjourney_v7_draft" => Ok(Self::MidjourneyV7Draft), "midjourney_v7_draft_raw" => Ok(Self::MidjourneyV7DraftRaw), + "midjourney_v7_raw" => Ok(Self::MidjourneyV7Raw), // Video models "kling_1p6_pro" => Ok(Self::Kling16Pro), @@ -149,9 +161,12 @@ impl ModelType { Self::Recraft3, Self::Midjourney, Self::MidjourneyV6, + Self::MidjourneyV6p1, + Self::MidjourneyV6p1Raw, Self::MidjourneyV7, - Self::MidjourneyV7Raw, + Self::MidjourneyV7Draft, Self::MidjourneyV7DraftRaw, + Self::MidjourneyV7Raw, // Video models Self::Kling16Pro, @@ -188,9 +203,12 @@ mod tests { assert_serialization(ModelType::Recraft3, "recraft_3"); assert_serialization(ModelType::Midjourney, "midjourney"); assert_serialization(ModelType::MidjourneyV6, "midjourney_v6"); + assert_serialization(ModelType::MidjourneyV6p1, "midjourney_v6p1"); + assert_serialization(ModelType::MidjourneyV6p1Raw, "midjourney_v6p1_raw"); assert_serialization(ModelType::MidjourneyV7, "midjourney_v7"); - assert_serialization(ModelType::MidjourneyV7Raw, "midjourney_v7_raw"); + assert_serialization(ModelType::MidjourneyV7Draft, "midjourney_v7_draft"); assert_serialization(ModelType::MidjourneyV7DraftRaw, "midjourney_v7_draft_raw"); + assert_serialization(ModelType::MidjourneyV7Raw, "midjourney_v7_raw"); // Video models assert_serialization(ModelType::Kling16Pro, "kling_1p6_pro"); assert_serialization(ModelType::Kling21Pro, "kling_2p1_pro"); @@ -215,9 +233,12 @@ mod tests { assert_eq!(ModelType::Recraft3.to_str(), "recraft_3"); assert_eq!(ModelType::Midjourney.to_str(), "midjourney"); assert_eq!(ModelType::MidjourneyV6.to_str(), "midjourney_v6"); + assert_eq!(ModelType::MidjourneyV6p1.to_str(), "midjourney_v6p1"); + assert_eq!(ModelType::MidjourneyV6p1Raw.to_str(), "midjourney_v6p1_raw"); assert_eq!(ModelType::MidjourneyV7.to_str(), "midjourney_v7"); - assert_eq!(ModelType::MidjourneyV7Raw.to_str(), "midjourney_v7_raw"); + assert_eq!(ModelType::MidjourneyV7Draft.to_str(), "midjourney_v7_draft"); assert_eq!(ModelType::MidjourneyV7DraftRaw.to_str(), "midjourney_v7_draft_raw"); + assert_eq!(ModelType::MidjourneyV7Raw.to_str(), "midjourney_v7_raw"); // Video models assert_eq!(ModelType::Kling16Pro.to_str(), "kling_1p6_pro"); @@ -244,9 +265,12 @@ mod tests { assert_eq!(ModelType::from_str("recraft_3").unwrap(), ModelType::Recraft3); assert_eq!(ModelType::from_str("midjourney").unwrap(), ModelType::Midjourney); assert_eq!(ModelType::from_str("midjourney_v6").unwrap(), ModelType::MidjourneyV6); + assert_eq!(ModelType::from_str("midjourney_v6p1").unwrap(), ModelType::MidjourneyV6p1); + assert_eq!(ModelType::from_str("midjourney_v6p1_raw").unwrap(), ModelType::MidjourneyV6p1Raw); assert_eq!(ModelType::from_str("midjourney_v7").unwrap(), ModelType::MidjourneyV7); - assert_eq!(ModelType::from_str("midjourney_v7_raw").unwrap(), ModelType::MidjourneyV7Raw); + assert_eq!(ModelType::from_str("midjourney_v7_draft").unwrap(), ModelType::MidjourneyV7Draft); assert_eq!(ModelType::from_str("midjourney_v7_draft_raw").unwrap(), ModelType::MidjourneyV7DraftRaw); + assert_eq!(ModelType::from_str("midjourney_v7_raw").unwrap(), ModelType::MidjourneyV7Raw); // Video models assert_eq!(ModelType::from_str("kling_1p6_pro").unwrap(), ModelType::Kling16Pro); assert_eq!(ModelType::from_str("kling_2p1_pro").unwrap(), ModelType::Kling21Pro); @@ -261,7 +285,7 @@ mod tests { #[test] fn all_variants() { let mut variants = ModelType::all_variants(); - assert_eq!(variants.len(), 20); + assert_eq!(variants.len(), 23); // Image models assert_eq!(variants.pop_first(), Some(ModelType::Flux1Dev)); assert_eq!(variants.pop_first(), Some(ModelType::Flux1Schnell)); @@ -273,9 +297,12 @@ mod tests { assert_eq!(variants.pop_first(), Some(ModelType::Recraft3)); assert_eq!(variants.pop_first(), Some(ModelType::Midjourney)); assert_eq!(variants.pop_first(), Some(ModelType::MidjourneyV6)); + assert_eq!(variants.pop_first(), Some(ModelType::MidjourneyV6p1)); + assert_eq!(variants.pop_first(), Some(ModelType::MidjourneyV6p1Raw)); assert_eq!(variants.pop_first(), Some(ModelType::MidjourneyV7)); - assert_eq!(variants.pop_first(), Some(ModelType::MidjourneyV7Raw)); + assert_eq!(variants.pop_first(), Some(ModelType::MidjourneyV7Draft)); assert_eq!(variants.pop_first(), Some(ModelType::MidjourneyV7DraftRaw)); + assert_eq!(variants.pop_first(), Some(ModelType::MidjourneyV7Raw)); // Video models assert_eq!(variants.pop_first(), Some(ModelType::Kling16Pro)); assert_eq!(variants.pop_first(), Some(ModelType::Kling21Pro)); diff --git a/frontend/libs/model-list/src/lib/ModelMapping.ts b/frontend/libs/model-list/src/lib/ModelMapping.ts index fae9c5403c..b6c1c3aa80 100644 --- a/frontend/libs/model-list/src/lib/ModelMapping.ts +++ b/frontend/libs/model-list/src/lib/ModelMapping.ts @@ -17,6 +17,12 @@ export const MODEL_TYPE_TO_CREATOR: Record = { recraft_3: ModelCreator.Recraft, hunyuan_3d: ModelCreator.Tencent, midjourney: ModelCreator.Midjourney, + midjourney_v6: ModelCreator.Midjourney, + midjourney_v6p1: ModelCreator.Midjourney, + midjourney_v6p1_raw: ModelCreator.Midjourney, + midjourney_v7: ModelCreator.Midjourney, + midjourney_v7_raw: ModelCreator.Midjourney, + midjourney_v7_draft_raw: ModelCreator.Midjourney, }; // Get creator icon for a model type @@ -95,7 +101,18 @@ export const getModelDisplayName = (modelType: string): string => { hunyuan_3d_2p1: "Hunyuan 3D 2.1", hunyuan_3d: "Hunyuan 3D", flux_pro_kontext_max: "Flux Pro Kontext Max", + + // Catch-all bucket for Midjourney. midjourney: "Midjourney", + + // Specific Midjourney models. + midjourney_v6: "Midjourney V6", + midjourney_v6p1: "Midjourney V6.1", + midjourney_v6p1_raw: "Midjourney V6.1 (Raw)", + midjourney_v7: "Midjourney V7", + midjourney_v7_raw: "Midjourney V7 (Raw)", + midjourney_v7_draft_raw: "Midjourney V7 (Draft Raw)", + // TODO: add more models here - BFlat };