diff --git a/.sqlx/query-0848cd9ddf8f61f2dd4d57f0c58b66a60a985acaa1ee5753f30f5e6e107e080e.json b/.sqlx/query-0848cd9ddf8f61f2dd4d57f0c58b66a60a985acaa1ee5753f30f5e6e107e080e.json deleted file mode 100644 index 3a77f3e997..0000000000 --- a/.sqlx/query-0848cd9ddf8f61f2dd4d57f0c58b66a60a985acaa1ee5753f30f5e6e107e080e.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "db_name": "SQLite", - "query": "\n INSERT INTO tasks (\n id,\n task_status,\n task_type,\n provider,\n provider_job_id,\n frontend_subscriber_id,\n frontend_subscriber_payload\n )\n VALUES (?, ?, ?, ?, ?, ?, ?)\n ", - "describe": { - "columns": [], - "parameters": { - "Right": 7 - }, - "nullable": [] - }, - "hash": "0848cd9ddf8f61f2dd4d57f0c58b66a60a985acaa1ee5753f30f5e6e107e080e" -} diff --git a/.sqlx/query-22b91a21cbb88abc9bc17b338cf2f9418e08e414667d44476623eb141ab17b87.json b/.sqlx/query-22b91a21cbb88abc9bc17b338cf2f9418e08e414667d44476623eb141ab17b87.json new file mode 100644 index 0000000000..b5174eb1d8 --- /dev/null +++ b/.sqlx/query-22b91a21cbb88abc9bc17b338cf2f9418e08e414667d44476623eb141ab17b87.json @@ -0,0 +1,12 @@ +{ + "db_name": "SQLite", + "query": "\n INSERT INTO tasks (\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n frontend_subscriber_id,\n frontend_subscriber_payload\n )\n VALUES (?, ?, ?, ?, ?, ?, ?, ?)\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 8 + }, + "nullable": [] + }, + "hash": "22b91a21cbb88abc9bc17b338cf2f9418e08e414667d44476623eb141ab17b87" +} diff --git a/_sql/artcraft_migrations/20250708000000_create_tasks_table.sql b/_sql/artcraft_migrations/20250708000000_create_tasks_table.sql index f2917def4b..9b0c4d0329 100644 --- a/_sql/artcraft_migrations/20250708000000_create_tasks_table.sql +++ b/_sql/artcraft_migrations/20250708000000_create_tasks_table.sql @@ -1,4 +1,7 @@ -- noinspection SqlDialectInspectionForFile +-- Migration history: +-- tasks_v1.sqlite - initial version +-- tasks_v2.sqlite - added model_type (nullable) CREATE TABLE tasks ( id TEXT NOT NULL PRIMARY KEY, @@ -6,6 +9,8 @@ CREATE TABLE tasks ( task_status TEXT NOT NULL, task_type TEXT NOT NULL, + model_type TEXT, + provider TEXT, provider_job_id TEXT, diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs index f2e71b52f6..7fd8d54686 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs @@ -1,6 +1,7 @@ use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider}; use crate::core::state::task_database::TaskDatabase; use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_model_type::TaskModelType; use enums::tauri::tasks::task_status::TaskStatus; use enums::tauri::tasks::task_type::TaskType; use sqlite_tasks::error::SqliteTasksError; @@ -34,10 +35,31 @@ impl TaskEnqueueSuccess{ } pub async fn insert_into_task_database(&self, task_database: &TaskDatabase) -> Result { + let model_type = match self.model { + None => None, + Some(GenerationModel::Flux1Dev) => Some(TaskModelType::Flux1Dev), + Some(GenerationModel::Flux1Schnell) => Some(TaskModelType::Flux1Schnell), + Some(GenerationModel::FluxPro11) => Some(TaskModelType::FluxPro11), + Some(GenerationModel::FluxPro11Ultra) => Some(TaskModelType::FluxPro11Ultra), + Some(GenerationModel::GptImage1) => Some(TaskModelType::GptImage1), + Some(GenerationModel::Recraft3) => Some(TaskModelType::Recraft3), + Some(GenerationModel::Kling21Pro) => Some(TaskModelType::Kling21Pro), + Some(GenerationModel::Kling21Master) => Some(TaskModelType::Kling21Master), + Some(GenerationModel::Seedance10Lite) => Some(TaskModelType::Seedance10Lite), + Some(GenerationModel::Veo2) => Some(TaskModelType::Veo2), + + Some(GenerationModel::Kling1_6) => Some(TaskModelType::Kling16Pro), // TODO: not accurate! + Some(GenerationModel::Hunyuan3d2_0) => None, // TODO + Some(GenerationModel::Hunyuan3d2_1) => None, // TODO + Some(GenerationModel::Kling2_0) => None, // TODO + Some(GenerationModel::Sora) => None, // TODO + }; + create_task(CreateTaskArgs { db: task_database.get_connection(), status: TaskStatus::Pending, task_type: self.task_type, + model_type: model_type, provider: self.provider, provider_job_id: self.provider_job_id.as_deref(), frontend_subscriber_id: None, diff --git a/crates/desktop/artcraft/src/core/state/data_dir/app_state_dir.rs b/crates/desktop/artcraft/src/core/state/data_dir/app_state_dir.rs index 9132e32d91..0af3d3f7d4 100644 --- a/crates/desktop/artcraft/src/core/state/data_dir/app_state_dir.rs +++ b/crates/desktop/artcraft/src/core/state/data_dir/app_state_dir.rs @@ -22,7 +22,7 @@ impl DataSubdir for AppStateDir{ impl AppStateDir { pub fn get_tasks_sqlite_database_path(&self) -> PathBuf { - self.path.join("tasks_v1.sqlite") + self.path.join("tasks_v2.sqlite") } pub fn get_window_size_config_file(&self) -> PathBuf { diff --git a/crates/schema/database/sqlite_tasks/src/queries/create_task.rs b/crates/schema/database/sqlite_tasks/src/queries/create_task.rs index 94323eeaf8..23f6888d34 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/create_task.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/create_task.rs @@ -1,6 +1,7 @@ use crate::connection::TaskDbConnection; use crate::error::SqliteTasksError; use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_model_type::TaskModelType; use enums::tauri::tasks::task_status::TaskStatus; use enums::tauri::tasks::task_type::TaskType; use tokens::tokens::sqlite::tasks::TaskId; @@ -9,6 +10,7 @@ pub struct CreateTaskArgs<'a> { pub db: &'a TaskDbConnection, pub status: TaskStatus, pub task_type: TaskType, + pub model_type: Option, pub provider: GenerationProvider, pub provider_job_id: Option<&'a str>, pub frontend_subscriber_id: Option<&'a str>, @@ -24,6 +26,7 @@ pub async fn create_task( let task_id_temp = task_id.as_str(); let status_temp = args.status.to_str(); let task_type_temp = args.task_type.to_str(); + let model_type_temp = args.model_type.map(|s| s.to_str()); let provider_temp = args.provider.to_string(); let query = sqlx::query!(r#" @@ -31,16 +34,18 @@ pub async fn create_task( id, task_status, task_type, + model_type, provider, provider_job_id, frontend_subscriber_id, frontend_subscriber_payload ) - VALUES (?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) "#, task_id_temp, status_temp, task_type_temp, + model_type_temp, provider_temp, args.provider_job_id, args.frontend_subscriber_id, diff --git a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs index bb60ee90ad..76284c579b 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs @@ -1,6 +1,7 @@ use crate::connection::TaskDbConnection; use crate::error::SqliteTasksError; use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_model_type::TaskModelType; use enums::tauri::tasks::task_status::TaskStatus; use enums::tauri::tasks::task_type::TaskType; use sqlx::{QueryBuilder, Sqlite}; @@ -21,6 +22,7 @@ pub struct Task { pub id: TaskId, pub status: TaskStatus, pub task_type: TaskType, + pub model_type: TaskModelType, pub provider: GenerationProvider, pub provider_job_id: Option, pub frontend_subscriber_id: Option, @@ -33,6 +35,7 @@ struct RawTask { id: String, task_status: String, task_type: String, + model_type: String, provider: String, provider_job_id: Option, frontend_subscriber_id: Option, @@ -48,6 +51,7 @@ pub async fn list_tasks_by_provider_and_tokens( id, task_status, task_type, + model_type, provider, provider_job_id, frontend_subscriber_id, @@ -87,6 +91,8 @@ pub async fn list_tasks_by_provider_and_tokens( .map_err(|err| SqliteTasksError::TaskParseError(err))?, task_type: TaskType::from_str(&task.task_type) .map_err(|err| SqliteTasksError::TaskParseError(err))?, + model_type: TaskModelType::from_str(&task.model_type) + .map_err(|err| SqliteTasksError::TaskParseError(err))?, provider: GenerationProvider::from_str(&task.provider) .map_err(|err| SqliteTasksError::TaskParseError(err))?, provider_job_id: task.provider_job_id, diff --git a/crates/schema/public/enums/src/tauri/tasks/mod.rs b/crates/schema/public/enums/src/tauri/tasks/mod.rs index af78d0cda9..352e6fdb3a 100644 --- a/crates/schema/public/enums/src/tauri/tasks/mod.rs +++ b/crates/schema/public/enums/src/tauri/tasks/mod.rs @@ -1,2 +1,3 @@ +pub mod task_model_type; pub mod task_status; pub mod task_type; \ No newline at end of file diff --git a/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs b/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs new file mode 100644 index 0000000000..dae0be269e --- /dev/null +++ b/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs @@ -0,0 +1,204 @@ +use std::collections::BTreeSet; + +#[cfg(test)] +use strum::EnumCount; +#[cfg(test)] +use strum::EnumIter; +use utoipa::ToSchema; + +#[cfg_attr(test, derive(EnumIter, EnumCount))] +#[derive(Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum TaskModelType { + // Image models + #[serde(rename = "flux_1_dev")] + Flux1Dev, + #[serde(rename = "flux_1_schnell")] + Flux1Schnell, + #[serde(rename = "flux_pro_1.1")] + FluxPro11, + #[serde(rename = "flux_pro_1.1_ultra")] + FluxPro11Ultra, + #[serde(rename = "gpt_image_1")] + GptImage1, + #[serde(rename = "recraft_3")] + Recraft3, + + // Video models + #[serde(rename = "kling_1.6_pro")] + Kling16Pro, + #[serde(rename = "kling_2.1_pro")] + Kling21Pro, + #[serde(rename = "kling_2.1_master")] + Kling21Master, + #[serde(rename = "seedance_1.0_lite")] + Seedance10Lite, + #[serde(rename = "veo_2")] + Veo2, +} + +impl_enum_display_and_debug_using_to_str!(TaskModelType); +//impl_mysql_enum_coders!(TaskModelType); +//impl_mysql_from_row!(TaskModelType); + +// NB: We can derive `sqlx::Type` instead of using `impl_mysql_enum_coders` + +impl TaskModelType { + pub fn to_str(&self) -> &'static str { + match self { + // Image models + Self::Flux1Dev => "flux_1_dev", + Self::Flux1Schnell => "flux_1_schnell", + Self::FluxPro11 => "flux_pro_1.1", + Self::FluxPro11Ultra => "flux_pro_1.1_ultra", + Self::GptImage1 => "gpt_image_1", + Self::Recraft3 => "recraft_3", + // Video models + Self::Kling16Pro => "kling_1.6_pro", + Self::Kling21Pro => "kling_2.1_pro", + Self::Kling21Master => "kling_2.1_master", + Self::Seedance10Lite => "seedance_1.0_lite", + Self::Veo2 => "veo_2", + } + } + + pub fn from_str(job_status: &str) -> Result { + match job_status { + // Image models + "flux_1_dev" => Ok(Self::Flux1Dev), + "flux_1_schnell" => Ok(Self::Flux1Schnell), + "flux_pro_1.1" => Ok(Self::FluxPro11), + "flux_pro_1.1_ultra" => Ok(Self::FluxPro11Ultra), + "gpt_image_1" => Ok(Self::GptImage1), + "recraft_3" => Ok(Self::Recraft3), + // Video models + "kling_1.6_pro" => Ok(Self::Kling16Pro), + "kling_2.1_pro" => Ok(Self::Kling21Pro), + "kling_2.1_master" => Ok(Self::Kling21Master), + "seedance_1.0_lite" => Ok(Self::Seedance10Lite), + "veo_2" => Ok(Self::Veo2), + _ => Err(format!("invalid task_state: {:?}", job_status)), + } + } + + pub fn all_variants() -> BTreeSet { + // NB: BTreeSet is sorted + // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. + BTreeSet::from([ + // 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, + ]) + } +} + +#[cfg(test)] +mod tests { + use crate::test_helpers::assert_serialization; + use crate::tauri::tasks::task_model_type::TaskModelType; + + mod explicit_checks { + use super::*; + + #[test] + fn test_serialization() { + // Image models + assert_serialization(TaskModelType::Flux1Dev, "flux_1_dev"); + assert_serialization(TaskModelType::Flux1Schnell, "flux_1_schnell"); + assert_serialization(TaskModelType::FluxPro11, "flux_pro_1.1"); + assert_serialization(TaskModelType::FluxPro11Ultra, "flux_pro_1.1_ultra"); + assert_serialization(TaskModelType::GptImage1, "gpt_image_1"); + assert_serialization(TaskModelType::Recraft3, "recraft_3"); + // Video models + assert_serialization(TaskModelType::Kling16Pro, "kling_1.6_pro"); + assert_serialization(TaskModelType::Kling21Pro, "kling_2.1_pro"); + assert_serialization(TaskModelType::Kling21Master, "kling_2.1_master"); + assert_serialization(TaskModelType::Seedance10Lite, "seedance_1.0_lite"); + assert_serialization(TaskModelType::Veo2, "veo_2"); + } + + #[test] + fn to_str() { + // Image models + assert_eq!(TaskModelType::Flux1Dev.to_str(), "flux_1_dev"); + assert_eq!(TaskModelType::Flux1Schnell.to_str(), "flux_1_schnell"); + assert_eq!(TaskModelType::FluxPro11.to_str(), "flux_pro_1.1"); + assert_eq!(TaskModelType::FluxPro11Ultra.to_str(), "flux_pro_1.1_ultra"); + assert_eq!(TaskModelType::GptImage1.to_str(), "gpt_image_1"); + assert_eq!(TaskModelType::Recraft3.to_str(), "recraft_3"); + // Video models + assert_eq!(TaskModelType::Kling16Pro.to_str(), "kling_1.6_pro"); + assert_eq!(TaskModelType::Kling21Pro.to_str(), "kling_2.1_pro"); + assert_eq!(TaskModelType::Kling21Master.to_str(), "kling_2.1_master"); + assert_eq!(TaskModelType::Seedance10Lite.to_str(), "seedance_1.0_lite"); + assert_eq!(TaskModelType::Veo2.to_str(), "veo_2"); + } + + #[test] + fn from_str() { + // Image models + assert_eq!(TaskModelType::from_str("flux_1_dev").unwrap(), TaskModelType::Flux1Dev); + assert_eq!(TaskModelType::from_str("flux_1_schnell").unwrap(), TaskModelType::Flux1Schnell); + assert_eq!(TaskModelType::from_str("flux_pro_1.1").unwrap(), TaskModelType::FluxPro11); + assert_eq!(TaskModelType::from_str("flux_pro_1.1_ultra").unwrap(), TaskModelType::FluxPro11Ultra); + assert_eq!(TaskModelType::from_str("gpt_image_1").unwrap(), TaskModelType::GptImage1); + assert_eq!(TaskModelType::from_str("recraft_3").unwrap(), TaskModelType::Recraft3); + // Video models + assert_eq!(TaskModelType::from_str("kling_1.6_pro").unwrap(), TaskModelType::Kling16Pro); + assert_eq!(TaskModelType::from_str("kling_2.1_pro").unwrap(), TaskModelType::Kling21Pro); + assert_eq!(TaskModelType::from_str("kling_2.1_master").unwrap(), TaskModelType::Kling21Master); + assert_eq!(TaskModelType::from_str("seedance_1.0_lite").unwrap(), TaskModelType::Seedance10Lite); + assert_eq!(TaskModelType::from_str("veo_2").unwrap(), TaskModelType::Veo2); + } + + #[test] + fn all_variants() { + let mut variants = TaskModelType::all_variants(); + assert_eq!(variants.len(), 11); + // Image models + assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Dev)); + assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Schnell)); + assert_eq!(variants.pop_first(), Some(TaskModelType::FluxPro11)); + assert_eq!(variants.pop_first(), Some(TaskModelType::FluxPro11Ultra)); + assert_eq!(variants.pop_first(), Some(TaskModelType::GptImage1)); + assert_eq!(variants.pop_first(), Some(TaskModelType::Recraft3)); + // Video models + assert_eq!(variants.pop_first(), Some(TaskModelType::Kling16Pro)); + assert_eq!(variants.pop_first(), Some(TaskModelType::Kling21Pro)); + assert_eq!(variants.pop_first(), Some(TaskModelType::Kling21Master)); + assert_eq!(variants.pop_first(), Some(TaskModelType::Seedance10Lite)); + assert_eq!(variants.pop_first(), Some(TaskModelType::Veo2)); + assert_eq!(variants.pop_first(), None); + } + } + + mod mechanical_checks { + use super::*; + + #[test] + fn variant_length() { + use strum::IntoEnumIterator; + assert_eq!(TaskModelType::all_variants().len(), TaskModelType::iter().len()); + } + + #[test] + fn round_trip() { + for variant in TaskModelType::all_variants() { + // Test to_str(), from_str(), Display, and Debug. + assert_eq!(variant, TaskModelType::from_str(variant.to_str()).unwrap()); + assert_eq!(variant, TaskModelType::from_str(&format!("{}", variant)).unwrap()); + assert_eq!(variant, TaskModelType::from_str(&format!("{:?}", variant)).unwrap()); + } + } + } +}