diff --git a/crates/desktop/artcraft/src/core/events/functional_events/mod.rs b/crates/desktop/artcraft/src/core/events/functional_events/mod.rs index 60db4d7b6c..605d0a82c2 100644 --- a/crates/desktop/artcraft/src/core/events/functional_events/mod.rs +++ b/crates/desktop/artcraft/src/core/events/functional_events/mod.rs @@ -1,8 +1,10 @@ pub mod canvas_background_removal_complete_event; pub mod credits_balance_changed_event; pub mod image_edit_complete_event; +pub mod object_generation_complete_event; pub mod refresh_account_state_event; pub mod show_provider_billing_modal_event; pub mod show_provider_login_modal_event; pub mod subscription_plan_changed_event; -pub mod text_to_image_generation_complete_event; \ No newline at end of file +pub mod text_to_image_generation_complete_event; +pub mod video_generation_complete_event; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/events/functional_events/object_generation_complete_event.rs b/crates/desktop/artcraft/src/core/events/functional_events/object_generation_complete_event.rs new file mode 100644 index 0000000000..990495270e --- /dev/null +++ b/crates/desktop/artcraft/src/core/events/functional_events/object_generation_complete_event.rs @@ -0,0 +1,26 @@ +use crate::core::events::basic_sendable_event_trait::{BasicEventStatus, BasicSendableEvent}; +use enums::tauri::ux::tauri_event_name::TauriEventName; +use serde_derive::Serialize; +use tokens::tokens::media_files::MediaFileToken; +use url::Url; + +#[derive(Clone, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +pub struct ObjectGenerationCompleteEvent { + pub generated_video: Option, + pub maybe_frontend_subscriber_id: Option, + pub maybe_frontend_subscriber_payload: Option, +} + +#[derive(Clone, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +pub struct GeneratedObject { + pub media_token: MediaFileToken, + pub cdn_url: Url, + pub maybe_thumbnail_template: Option, +} + +impl BasicSendableEvent for ObjectGenerationCompleteEvent { + const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::ObjectGenerationCompleteEvent; + const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Success; +} diff --git a/crates/desktop/artcraft/src/core/events/functional_events/video_generation_complete_event.rs b/crates/desktop/artcraft/src/core/events/functional_events/video_generation_complete_event.rs new file mode 100644 index 0000000000..5e20bcf548 --- /dev/null +++ b/crates/desktop/artcraft/src/core/events/functional_events/video_generation_complete_event.rs @@ -0,0 +1,26 @@ +use crate::core::events::basic_sendable_event_trait::{BasicEventStatus, BasicSendableEvent}; +use enums::tauri::ux::tauri_event_name::TauriEventName; +use serde_derive::Serialize; +use tokens::tokens::media_files::MediaFileToken; +use url::Url; + +#[derive(Clone, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +pub struct VideoGenerationCompleteEvent { + pub generated_video: Option, + pub maybe_frontend_subscriber_id: Option, + pub maybe_frontend_subscriber_payload: Option, +} + +#[derive(Clone, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +pub struct GeneratedVideo { + pub media_token: MediaFileToken, + pub cdn_url: Url, + pub maybe_thumbnail_template: Option, +} + +impl BasicSendableEvent for VideoGenerationCompleteEvent { + const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::VideoGenerationCompleteEvent; + const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Success; +} diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_frontend_caller_notification.rs b/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_frontend_caller_notification.rs new file mode 100644 index 0000000000..3798f6b686 --- /dev/null +++ b/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_frontend_caller_notification.rs @@ -0,0 +1,244 @@ +use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::functional_events::image_edit_complete_event::{EditedImage, ImageEditCompleteEvent}; +use crate::core::events::functional_events::object_generation_complete_event::{GeneratedObject, ObjectGenerationCompleteEvent}; +use crate::core::events::functional_events::text_to_image_generation_complete_event::{GeneratedImage, TextToImageGenerationCompleteEvent}; +use crate::core::events::functional_events::video_generation_complete_event::{GeneratedVideo, VideoGenerationCompleteEvent}; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use anyhow::anyhow; +use artcraft_api_defs::jobs::list_session_jobs::{ListSessionJobsItem, ListSessionResultDetailsResponse}; +use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template; +use enums::tauri::tasks::task_type::TaskType; +use enums::tauri::ux::tauri_command_caller::TauriCommandCaller; +use errors::AnyhowResult; +use log::{error, warn}; +use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs}; +use sqlite_tasks::queries::task::Task; +use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; +use storyteller_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files; +use tauri::AppHandle; +use tokens::tokens::batch_generations::BatchGenerationToken; +use tokens::tokens::media_files::MediaFileToken; +use crate::core::events::functional_events::canvas_background_removal_complete_event::CanvasBackgroundRemovalCompleteEvent; + +pub async fn maybe_handle_frontend_caller_notification( + app: &AppHandle, + app_env_configs: &AppEnvConfigs, + maybe_creds: Option<&StorytellerCredentialSet>, + task: &Task, + job: &ListSessionJobsItem, +) -> AnyhowResult<()> { + + let job_result = match job.maybe_result { + Some(ref res) => res, + None => { + warn!("Job result is None for task: {:?}", task); + return Ok(()); // No result, nothing to do + }, + }; + + match task.task_type { + TaskType::ImageGeneration => { + let _r = handle_image_generation( + app, + task, + job_result, + app_env_configs, + maybe_creds, + ).await?; + } + TaskType::ImageInpaintEdit => { + let _r = handle_inpaint_image_generation( + app, + task, + job_result, + app_env_configs, + maybe_creds, + ).await?; + } + TaskType::VideoGeneration => { + let _r = handle_video_generation( + app, + task, + job_result, + ).await?; + } + TaskType::ObjectGeneration => { + let _r = handle_object_generation( + app, + task, + job_result, + ).await?; + } + TaskType::BackgroundRemoval => { + let _r = handle_background_removal_generation( + app, + task, + job_result, + ).await?; + } + } + + Ok(()) +} + +async fn handle_image_generation( + app: &AppHandle, + task: &Task, + job_result: &ListSessionResultDetailsResponse, + app_env_configs: &AppEnvConfigs, + maybe_creds: Option<&StorytellerCredentialSet>, +) -> AnyhowResult<()> { + + let generated_images = match job_result.maybe_batch_token.as_ref() { + None => { + vec![GeneratedImage { + media_token: MediaFileToken::new_from_str(&job_result.entity_token), + cdn_url: job_result.media_links.cdn_url.clone(), + maybe_thumbnail_template: job_result.media_links.maybe_thumbnail_template.clone(), + }] + } + Some(batch_token) => { + let result = list_batch_generated_redux_media_files( + &app_env_configs.storyteller_host, + maybe_creds, + batch_token, + ).await?; + + if result.media_files.is_empty() { + return Err(anyhow!("No media files found for batch token: {}", batch_token)); + } + + result.media_files + .into_iter() + .map(|file| GeneratedImage { + media_token: file.token, + cdn_url: file.media_links.cdn_url, + maybe_thumbnail_template: file.media_links.maybe_thumbnail_template, + }) + .collect() + } + }; + + let event = TextToImageGenerationCompleteEvent { + generated_images, + maybe_frontend_subscriber_id: task.frontend_subscriber_id.clone(), + maybe_frontend_subscriber_payload: task.frontend_subscriber_payload.clone(), + }; + + event.send_infallible(&app); + + Ok(()) +} + +async fn handle_inpaint_image_generation( + app: &AppHandle, + task: &Task, + job_result: &ListSessionResultDetailsResponse, + app_env_configs: &AppEnvConfigs, + maybe_creds: Option<&StorytellerCredentialSet>, +) -> AnyhowResult<()> { + + let edited_images = match job_result.maybe_batch_token.as_ref() { + None => { + vec![EditedImage { + media_token: MediaFileToken::new_from_str(&job_result.entity_token), + cdn_url: job_result.media_links.cdn_url.clone(), + maybe_thumbnail_template: job_result.media_links.maybe_thumbnail_template.clone(), + }] + } + Some(batch_token) => { + let result = list_batch_generated_redux_media_files( + &app_env_configs.storyteller_host, + maybe_creds, + batch_token, + ).await?; + + if result.media_files.is_empty() { + return Err(anyhow!("No media files found for batch token: {}", batch_token)); + } + + result.media_files + .into_iter() + .map(|file| EditedImage { + media_token: file.token, + cdn_url: file.media_links.cdn_url, + maybe_thumbnail_template: file.media_links.maybe_thumbnail_template, + }) + .collect() + } + }; + + let event = ImageEditCompleteEvent { + edited_images, + maybe_frontend_subscriber_id: task.frontend_subscriber_id.clone(), + maybe_frontend_subscriber_payload: task.frontend_subscriber_payload.clone(), + }; + + event.send_infallible(&app); + + Ok(()) +} + +async fn handle_video_generation( + app: &AppHandle, + task: &Task, + job_result: &ListSessionResultDetailsResponse, +) -> AnyhowResult<()> { + + // NB: For now, we only generate one video at a time. + let event = VideoGenerationCompleteEvent { + generated_video: Some(GeneratedVideo { + media_token: MediaFileToken::new_from_str(&job_result.entity_token), + cdn_url: job_result.media_links.cdn_url.clone(), + maybe_thumbnail_template: media_links_to_thumbnail_template(&job_result.media_links) + .map(|s| s.to_owned()), + }), + maybe_frontend_subscriber_id: task.frontend_subscriber_id.clone(), + maybe_frontend_subscriber_payload: task.frontend_subscriber_payload.clone(), + }; + + event.send_infallible(&app); + + Ok(()) +} + +async fn handle_object_generation( + app: &AppHandle, + task: &Task, + job_result: &ListSessionResultDetailsResponse, +) -> AnyhowResult<()> { + + // NB: For now, we only generate one object (3d mesh) at a time. + let event = ObjectGenerationCompleteEvent { + generated_video: Some(GeneratedObject { + media_token: MediaFileToken::new_from_str(&job_result.entity_token), + cdn_url: job_result.media_links.cdn_url.clone(), + maybe_thumbnail_template: media_links_to_thumbnail_template(&job_result.media_links) + .map(|s| s.to_owned()), + }), + maybe_frontend_subscriber_id: task.frontend_subscriber_id.clone(), + maybe_frontend_subscriber_payload: task.frontend_subscriber_payload.clone(), + }; + + event.send_infallible(&app); + + Ok(()) +} + +async fn handle_background_removal_generation( + app: &AppHandle, + task: &Task, + job_result: &ListSessionResultDetailsResponse, +) -> AnyhowResult<()> { + + let event = CanvasBackgroundRemovalCompleteEvent { + media_token: MediaFileToken::new_from_str(&job_result.entity_token), + image_cdn_url: job_result.media_links.cdn_url.clone(), + maybe_frontend_subscriber_id: task.frontend_subscriber_id.clone(), + maybe_frontend_subscriber_payload: task.frontend_subscriber_payload.clone(), + }; + + event.send_infallible(&app); + + Ok(()) +} diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/events/mod.rs b/crates/desktop/artcraft/src/services/storyteller/threads/events/mod.rs index fa576b10ce..c9ae78ea32 100644 --- a/crates/desktop/artcraft/src/services/storyteller/threads/events/mod.rs +++ b/crates/desktop/artcraft/src/services/storyteller/threads/events/mod.rs @@ -1,3 +1,4 @@ +pub (super) mod maybe_handle_frontend_caller_notification; pub (super) mod maybe_handle_inpainting_complete_event; pub (super) mod maybe_handle_text_to_image_complete_event; pub (super) mod maybe_send_background_removal_complete_event; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs index 3f9c8f7f63..7c67fb8a2c 100644 --- a/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs +++ b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs @@ -12,6 +12,7 @@ use crate::core::utils::enum_conversion::task_type::to_generation_action; use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; use crate::services::sora::state::sora_task_queue::SoraTaskQueue; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use crate::services::storyteller::threads::events::maybe_handle_frontend_caller_notification::maybe_handle_frontend_caller_notification; use crate::services::storyteller::threads::events::maybe_handle_inpainting_complete_event::maybe_handle_inpainting_complete_event; use crate::services::storyteller::threads::events::maybe_handle_text_to_image_complete_event::maybe_handle_text_to_image_complete_event; use crate::services::storyteller::threads::events::maybe_send_background_removal_complete_event::maybe_send_background_removal_complete_event; @@ -197,7 +198,7 @@ async fn send_additional_events( ) { info!("Attempting to dispatch events for completed Storyteller : {:?}", task); - let result = maybe_handle_text_to_image_complete_event( + let result = maybe_handle_frontend_caller_notification( app_handle, app_env_configs, creds, @@ -206,30 +207,42 @@ async fn send_additional_events( ).await; if let Err(err) = result { - error!("Failed to send text-to-image complete event: {:?}", err); + error!("Failed to send generation complete event: {:?}", err); } - let result = maybe_handle_inpainting_complete_event( - app_handle, - app_env_configs, - creds, - task, - job, - ).await; + //let result = maybe_handle_text_to_image_complete_event( + // app_handle, + // app_env_configs, + // creds, + // task, + // job, + //).await; - if let Err(err) = result { - error!("Failed to send image inpainting complete event: {:?}", err); - } + //if let Err(err) = result { + // error!("Failed to send text-to-image complete event: {:?}", err); + //} - let result = maybe_send_background_removal_complete_event( - app_handle, - task, - job, - ).await; + //let result = maybe_handle_inpainting_complete_event( + // app_handle, + // app_env_configs, + // creds, + // task, + // job, + //).await; - if let Err(err) = result { - error!("Failed to send background removal complete event: {:?}", err); - } + //if let Err(err) = result { + // error!("Failed to send image inpainting complete event: {:?}", err); + //} + + //let result = maybe_send_background_removal_complete_event( + // app_handle, + // task, + // job, + //).await; + + //if let Err(err) = result { + // error!("Failed to send background removal complete event: {:?}", err); + //} } fn get_thumbnail_template<'a>(job: &'a ListSessionJobsItem) -> Option<&'a str> { diff --git a/crates/schema/public/enums/src/tauri/ux/tauri_event_name.rs b/crates/schema/public/enums/src/tauri/ux/tauri_event_name.rs index fd6b8bf44b..3c61656bfb 100644 --- a/crates/schema/public/enums/src/tauri/ux/tauri_event_name.rs +++ b/crates/schema/public/enums/src/tauri/ux/tauri_event_name.rs @@ -54,7 +54,7 @@ pub enum TauriEventName { CanvasBgRemovedEvent, /// Special event: - /// Image edit is complete + /// Image generation is complete #[serde(rename = "text_to_image_generation_complete_event")] TextToImageGenerationCompleteEvent, @@ -62,6 +62,16 @@ pub enum TauriEventName { /// Image edit is complete #[serde(rename = "image_edit_complete_event")] ImageEditCompleteEvent, + + /// Special event: + /// Object (3D mesh) generation is complete + #[serde(rename = "object_generation_complete_event")] + ObjectGenerationCompleteEvent, + + /// Special event: + /// Video generation is complete + #[serde(rename = "video_generation_complete_event")] + VideoGenerationCompleteEvent, /// Special event: /// Refresh account states @@ -102,6 +112,8 @@ impl TauriEventName { Self::CanvasBgRemovedEvent => "canvas_bg_removed_event", Self::TextToImageGenerationCompleteEvent => "text_to_image_generation_complete_event", Self::ImageEditCompleteEvent => "image_edit_complete_event", + Self::ObjectGenerationCompleteEvent => "object_generation_complete_event", + Self::VideoGenerationCompleteEvent => "video_generation_complete_event", Self::RefreshAccountStateEvent => "refresh_account_state_event", Self::ShowProviderBillingModalEvent => "show_provider_billing_modal_event", Self::ShowProviderLoginModalEvent => "show_provider_login_modal_event", @@ -120,6 +132,8 @@ impl TauriEventName { "canvas_bg_removed_event" => Ok(Self::CanvasBgRemovedEvent), "text_to_image_generation_complete_event" => Ok(Self::TextToImageGenerationCompleteEvent), "image_edit_complete_event" => Ok(Self::ImageEditCompleteEvent), + "object_generation_complete_event" => Ok(Self::ObjectGenerationCompleteEvent), + "video_generation_complete_event" => Ok(Self::VideoGenerationCompleteEvent), "refresh_account_state_event" => Ok(Self::RefreshAccountStateEvent), "show_provider_billing_modal_event" => Ok(Self::ShowProviderBillingModalEvent), "show_provider_login_modal_event" => Ok(Self::ShowProviderLoginModalEvent), @@ -141,6 +155,8 @@ impl TauriEventName { Self::CanvasBgRemovedEvent, Self::TextToImageGenerationCompleteEvent, Self::ImageEditCompleteEvent, + Self::ObjectGenerationCompleteEvent, + Self::VideoGenerationCompleteEvent, Self::RefreshAccountStateEvent, Self::ShowProviderBillingModalEvent, Self::ShowProviderLoginModalEvent, @@ -168,6 +184,8 @@ mod tests { assert_serialization(TauriEventName::CanvasBgRemovedEvent, "canvas_bg_removed_event"); assert_serialization(TauriEventName::TextToImageGenerationCompleteEvent, "text_to_image_generation_complete_event"); assert_serialization(TauriEventName::ImageEditCompleteEvent, "image_edit_complete_event"); + assert_serialization(TauriEventName::ObjectGenerationCompleteEvent, "object_generation_complete_event"); + assert_serialization(TauriEventName::VideoGenerationCompleteEvent, "video_generation_complete_event"); assert_serialization(TauriEventName::RefreshAccountStateEvent, "refresh_account_state_event"); assert_serialization(TauriEventName::ShowProviderBillingModalEvent, "show_provider_billing_modal_event"); assert_serialization(TauriEventName::ShowProviderLoginModalEvent, "show_provider_login_modal_event"); @@ -185,6 +203,8 @@ mod tests { assert_eq!(TauriEventName::CanvasBgRemovedEvent.to_str(), "canvas_bg_removed_event"); assert_eq!(TauriEventName::TextToImageGenerationCompleteEvent.to_str(), "text_to_image_generation_complete_event"); assert_eq!(TauriEventName::ImageEditCompleteEvent.to_str(), "image_edit_complete_event"); + assert_eq!(TauriEventName::ObjectGenerationCompleteEvent.to_str(), "object_generation_complete_event"); + assert_eq!(TauriEventName::VideoGenerationCompleteEvent.to_str(), "video_generation_complete_event"); assert_eq!(TauriEventName::RefreshAccountStateEvent.to_str(), "refresh_account_state_event"); assert_eq!(TauriEventName::ShowProviderBillingModalEvent.to_str(), "show_provider_billing_modal_event"); assert_eq!(TauriEventName::ShowProviderLoginModalEvent.to_str(), "show_provider_login_modal_event"); @@ -202,6 +222,8 @@ mod tests { assert_eq!(TauriEventName::from_str("canvas_bg_removed_event").unwrap(), TauriEventName::CanvasBgRemovedEvent); assert_eq!(TauriEventName::from_str("text_to_image_generation_complete_event").unwrap(), TauriEventName::TextToImageGenerationCompleteEvent); assert_eq!(TauriEventName::from_str("image_edit_complete_event").unwrap(), TauriEventName::ImageEditCompleteEvent); + assert_eq!(TauriEventName::from_str("object_generation_complete_event").unwrap(), TauriEventName::ObjectGenerationCompleteEvent); + assert_eq!(TauriEventName::from_str("video_generation_complete_event").unwrap(), TauriEventName::VideoGenerationCompleteEvent); assert_eq!(TauriEventName::from_str("refresh_account_state_event").unwrap(), TauriEventName::RefreshAccountStateEvent); assert_eq!(TauriEventName::from_str("show_provider_billing_modal_event").unwrap(), TauriEventName::ShowProviderBillingModalEvent); assert_eq!(TauriEventName::from_str("show_provider_login_modal_event").unwrap(), TauriEventName::ShowProviderLoginModalEvent); @@ -211,7 +233,7 @@ mod tests { #[test] fn all_variants() { let mut variants = TauriEventName::all_variants(); - assert_eq!(variants.len(), 13); + assert_eq!(variants.len(), 15); assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationEnqueueSuccessEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationEnqueueFailureEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationCompleteEvent)); @@ -221,6 +243,8 @@ mod tests { assert_eq!(variants.pop_first(), Some(TauriEventName::CanvasBgRemovedEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::TextToImageGenerationCompleteEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::ImageEditCompleteEvent)); + assert_eq!(variants.pop_first(), Some(TauriEventName::ObjectGenerationCompleteEvent)); + assert_eq!(variants.pop_first(), Some(TauriEventName::VideoGenerationCompleteEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::RefreshAccountStateEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::ShowProviderBillingModalEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::ShowProviderLoginModalEvent));