diff --git a/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/events/notify_frontend_of_completion.rs b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/events/notify_frontend_of_completion.rs index 5c366d03e6..324b71f67a 100644 --- a/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/events/notify_frontend_of_completion.rs +++ b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/events/notify_frontend_of_completion.rs @@ -1,4 +1,6 @@ use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider}; +use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent; use crate::core::events::functional_events::canvas_background_removal_complete_event::CanvasBackgroundRemovalCompleteEvent; use crate::core::events::functional_events::gaussian_generation_complete_event::{GaussianGenerationCompleteEvent, GeneratedGaussian}; use crate::core::events::functional_events::image_edit_complete_event::{EditedImage, ImageEditCompleteEvent}; @@ -10,7 +12,9 @@ use artcraft_client::endpoints::media_files::get_media_file::get_media_file; use artcraft_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files; use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet; use artcraft_client::utils::api_host::ApiHost; +use enums::common::generation_provider::GenerationProvider; use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass; +use enums::tauri::tasks::task_model_type::TaskModelType; use enums::tauri::tasks::task_type::TaskType; use log::{error, info, warn}; use sqlite_tasks::queries::task::Task; @@ -49,6 +53,19 @@ pub async fn notify_frontend_of_completion( task: &Task, completion: &CompletionData, ) { + // Fire the generic generation-complete event (for the task queue UI). + let generation_action = task_type_to_generation_action(task.task_type); + let generation_model = task.model_type.and_then(task_model_type_to_generation_model); + let generation_service = provider_to_generation_service(task.provider); + + let complete_event = GenerationCompleteEvent { + action: Some(generation_action), + service: generation_service, + model: generation_model, + }; + complete_event.send_infallible(app); + + // Fire the typed frontend notification (for the specific page/component that initiated the job). let result = match task.task_type { TaskType::ImageGeneration => { notify_image_generation(app, api_host, maybe_creds, task, completion).await @@ -237,3 +254,78 @@ async fn collect_generated_images( maybe_thumbnail_template: completion.maybe_thumbnail_url_template.clone(), }] } + +fn task_type_to_generation_action(task_type: TaskType) -> GenerationAction { + match task_type { + TaskType::ImageGeneration => GenerationAction::GenerateImage, + TaskType::VideoGeneration => GenerationAction::GenerateVideo, + TaskType::BackgroundRemoval => GenerationAction::RemoveBackground, + TaskType::ObjectGeneration => GenerationAction::ImageTo3d, + TaskType::GaussianGeneration => GenerationAction::GenerateGaussian, + TaskType::ImageInpaintEdit => GenerationAction::ImageInpaintEdit, + } +} + +fn provider_to_generation_service(provider: GenerationProvider) -> GenerationServiceProvider { + match provider { + GenerationProvider::Artcraft => GenerationServiceProvider::Artcraft, + GenerationProvider::Fal => GenerationServiceProvider::Fal, + GenerationProvider::Grok => GenerationServiceProvider::Grok, + GenerationProvider::Midjourney => GenerationServiceProvider::Midjourney, + GenerationProvider::Sora => GenerationServiceProvider::Sora, + GenerationProvider::WorldLabs => GenerationServiceProvider::WorldLabs, + } +} + +fn task_model_type_to_generation_model(model: TaskModelType) -> Option { + match model { + TaskModelType::Flux1Dev => Some(GenerationModel::Flux1Dev), + TaskModelType::Flux1Schnell => Some(GenerationModel::Flux1Schnell), + TaskModelType::FluxPro1 => Some(GenerationModel::FluxPro1), + TaskModelType::FluxPro11 => Some(GenerationModel::FluxPro11), + TaskModelType::FluxPro11Ultra => Some(GenerationModel::FluxPro11Ultra), + TaskModelType::FluxProKontextMax => Some(GenerationModel::FluxProKontextMax), + TaskModelType::FluxDevJuggernaut => Some(GenerationModel::FluxDevJuggernaut), + TaskModelType::Flux2LoraAngles => Some(GenerationModel::Flux2LoraAngles), + TaskModelType::GptImage1 => Some(GenerationModel::GptImage1), + TaskModelType::GptImage1p5 => Some(GenerationModel::GptImage1p5), + TaskModelType::GptImage2 => Some(GenerationModel::GptImage2), + TaskModelType::NanoBanana => Some(GenerationModel::NanoBanana), + TaskModelType::NanoBanana2 => Some(GenerationModel::NanoBanana2), + TaskModelType::NanoBananaPro => Some(GenerationModel::NanoBananaPro), + TaskModelType::Seedream4 => Some(GenerationModel::Seedream4), + TaskModelType::Seedream4p5 => Some(GenerationModel::Seedream4p5), + TaskModelType::Seedream5Lite => Some(GenerationModel::Seedream5Lite), + TaskModelType::QwenEdit2511Angles => Some(GenerationModel::QwenEdit2511Angles), + TaskModelType::GrokImage => Some(GenerationModel::GrokImage), + TaskModelType::Recraft3 => Some(GenerationModel::Recraft3), + TaskModelType::GrokVideo => Some(GenerationModel::GrokVideo), + TaskModelType::Kling16Pro => Some(GenerationModel::Kling1_6), + TaskModelType::Kling21Pro => Some(GenerationModel::Kling21Pro), + TaskModelType::Kling21Master => Some(GenerationModel::Kling21Master), + TaskModelType::Kling2p5TurboPro => Some(GenerationModel::Kling2p5TurboPro), + TaskModelType::Kling2p6Pro => Some(GenerationModel::Kling2p6Pro), + TaskModelType::Kling3p0Standard => Some(GenerationModel::Kling3p0Standard), + TaskModelType::Kling3p0Pro => Some(GenerationModel::Kling3p0Pro), + TaskModelType::HappyHorse1p0 => Some(GenerationModel::HappyHorse1p0), + TaskModelType::Seedance10Lite => Some(GenerationModel::Seedance10Lite), + TaskModelType::Seedance1p5Pro => Some(GenerationModel::Seedance1p5Pro), + TaskModelType::Seedance2p0 => Some(GenerationModel::Seedance2p0), + TaskModelType::Seedance2p0Fast => Some(GenerationModel::Seedance2p0Fast), + TaskModelType::Sora2 => Some(GenerationModel::Sora2), + TaskModelType::Sora2Pro => Some(GenerationModel::Sora2Pro), + TaskModelType::Veo2 => Some(GenerationModel::Veo2), + TaskModelType::Veo3 => Some(GenerationModel::Veo3), + TaskModelType::Veo3Fast => Some(GenerationModel::Veo3Fast), + TaskModelType::Veo3p1 => Some(GenerationModel::Veo3p1), + TaskModelType::Veo3p1Fast => Some(GenerationModel::Veo3p1Fast), + TaskModelType::Hunyuan3d2_0 => Some(GenerationModel::Hunyuan3d2_0), + TaskModelType::Hunyuan3d2_1 => Some(GenerationModel::Hunyuan3d2_1), + TaskModelType::Hunyuan3d3 => Some(GenerationModel::Hunyuan3d3), + TaskModelType::Midjourney => Some(GenerationModel::Midjourney), + TaskModelType::Gemini25Flash => Some(GenerationModel::Gemini25Flash), + TaskModelType::WorldlabsMarble => Some(GenerationModel::WorldlabsMarble), + TaskModelType::WorldlabsMarble0p1Mini => Some(GenerationModel::WorldlabsMarble0p1Mini), + TaskModelType::WorldlabsMarble0p1Plus => Some(GenerationModel::WorldlabsMarble0p1Plus), + } +} diff --git a/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_complete.rs b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_complete.rs index b40e3690bc..de3ddd51fe 100644 --- a/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_complete.rs +++ b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_complete.rs @@ -1,6 +1,3 @@ -use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; -use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; -use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent; use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::threads::third_party_task_polling_thread::events::notify_frontend_of_completion::{ notify_frontend_of_completion, CompletionData, @@ -84,7 +81,7 @@ async fn handle_fal_complete_inner( let extracted = job_response.extracted_contents; // Determine what kind of media we got and collect download URLs. - let (urls, media_class, generation_action) = collect_media_urls(task, &extracted)?; + let (urls, media_class) = collect_media_urls(task, &extracted)?; if urls.is_empty() { warn!("[FalComplete] Task {} completed but no downloadable media found in response", task.id.as_str()); @@ -166,15 +163,6 @@ async fn handle_fal_complete_inner( }).await?; if updated { - // Fire the generic generation-complete event (for the task queue UI) - let event = GenerationCompleteEvent { - action: Some(generation_action), - service: GenerationServiceProvider::Fal, - model: None, - }; - event.send_infallible(app_handle); - - // Fire the typed frontend notification (for the specific page/component that initiated the job) if let Some(primary_token) = maybe_primary_media_file_token { let completion = CompletionData { primary_media_file_token: primary_token, @@ -205,10 +193,10 @@ async fn handle_fal_complete_inner( fn collect_media_urls( task: &Task, extracted: &Option, -) -> Result<(Vec, TaskMediaFileClass, GenerationAction), Box> { +) -> Result<(Vec, TaskMediaFileClass), Box> { let extracted = match extracted { Some(e) => e, - None => return Ok((vec![], TaskMediaFileClass::Image, GenerationAction::GenerateImage)), + None => return Ok((vec![], TaskMediaFileClass::Image)), }; // Images (batch) @@ -217,32 +205,32 @@ fn collect_media_urls( .filter_map(|img| img.url.clone()) .collect(); if !urls.is_empty() { - return Ok((urls, TaskMediaFileClass::Image, GenerationAction::GenerateImage)); + return Ok((urls, TaskMediaFileClass::Image)); } } // Single image (e.g. background removal) if let Some(image) = &extracted.image { if let Some(url) = &image.url { - return Ok((vec![url.clone()], TaskMediaFileClass::Image, GenerationAction::GenerateImage)); + return Ok((vec![url.clone()], TaskMediaFileClass::Image)); } } // Video if let Some(video) = &extracted.video { if let Some(url) = &video.url { - return Ok((vec![url.clone()], TaskMediaFileClass::Video, GenerationAction::GenerateVideo)); + return Ok((vec![url.clone()], TaskMediaFileClass::Video)); } } // 3D model (GLB) if let Some(glb) = &extracted.model_glb { if let Some(url) = &glb.url { - return Ok((vec![url.clone()], TaskMediaFileClass::Dimensional, GenerationAction::ImageTo3d)); + return Ok((vec![url.clone()], TaskMediaFileClass::Dimensional)); } } - Ok((vec![], TaskMediaFileClass::Image, GenerationAction::GenerateImage)) + Ok((vec![], TaskMediaFileClass::Image)) } async fn download_file( @@ -353,3 +341,4 @@ async fn try_upload( Ok(media_token) } +