diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_bg_removal/enqueue_image_bg_removal_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_bg_removal/enqueue_image_bg_removal_command.rs index c195f65ca4..d4d3286a50 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_bg_removal/enqueue_image_bg_removal_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_bg_removal/enqueue_image_bg_removal_command.rs @@ -12,6 +12,8 @@ use crate::core::events::generation_events::generation_enqueue_failure_event::Ge use crate::core::events::generation_events::generation_enqueue_success_event::GenerationEnqueueSuccessEvent; use crate::core::events::sendable_event_trait::SendableEvent; use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_tracker::ArtcraftUsageTracker; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_type::{ArtcraftUsagePage, ArtcraftUsageType}; use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::core::state::data_dir::trait_data_subdir::DataSubdir; use crate::core::state::provider_priority::{Provider, ProviderPriorityStore}; @@ -97,6 +99,7 @@ pub async fn enqueue_image_bg_removal_command( request: EnqueueImageBgRemovalCommand, app_data_root: State<'_, AppDataRoot>, app_env_configs: State<'_, AppEnvConfigs>, + artcraft_usage_tracker: State<'_, ArtcraftUsageTracker>, provider_priority_store: State<'_, ProviderPriorityStore>, task_database: State<'_, TaskDatabase>, storyteller_creds_manager: State<'_, StorytellerCredentialManager>, @@ -109,6 +112,7 @@ pub async fn enqueue_image_bg_removal_command( &app, &app_data_root, &app_env_configs, + &artcraft_usage_tracker, &provider_priority_store, &task_database, &storyteller_creds_manager, @@ -157,6 +161,7 @@ pub async fn handle_request( app: &AppHandle, app_data_root: &AppDataRoot, app_env_configs: &AppEnvConfigs, + artcraft_usage_tracker: &ArtcraftUsageTracker, provider_priority_store: &ProviderPriorityStore, task_database: &TaskDatabase, storyteller_creds_manager: &StorytellerCredentialManager, @@ -185,6 +190,11 @@ pub async fn handle_request( error!("Failed to create task in database: {:?}", err); // NB: Fail open, but find a way to flag this. } + + if let Err(err) = artcraft_usage_tracker.record_image_generation(1, ArtcraftUsageType::ImageToResult, ArtcraftUsagePage::EditPage) { + // NB: Fail open. + warn!("Failed to report usage: {:?}", err); + } Ok(success_event) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/enqueue_edit_image_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/enqueue_edit_image_command.rs index d740ec276b..53f6cfbd72 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/enqueue_edit_image_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_edit/enqueue_edit_image_command.rs @@ -43,6 +43,8 @@ use storyteller_client::error::storyteller_error::StorytellerError; use storyteller_client::utils::api_host::ApiHost; use tauri::{AppHandle, Manager, State}; use tokens::tokens::media_files::MediaFileToken; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_tracker::ArtcraftUsageTracker; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_type::{ArtcraftUsagePage, ArtcraftUsageType}; /// This is used in the Tauri command bridge. /// Don't change the serializations without coordinating with the frontend. @@ -204,6 +206,7 @@ pub async fn enqueue_edit_image_command( request: EnqueueEditImageCommand, app_data_root: State<'_, AppDataRoot>, app_env_configs: State<'_, AppEnvConfigs>, + artcraft_usage_tracker: State<'_, ArtcraftUsageTracker>, provider_priority_store: State<'_, ProviderPriorityStore>, task_database: State<'_, TaskDatabase>, storyteller_creds_manager: State<'_, StorytellerCredentialManager>, @@ -219,6 +222,7 @@ pub async fn enqueue_edit_image_command( &app, &app_data_root, &app_env_configs, + &artcraft_usage_tracker, &provider_priority_store, &task_database, &storyteller_creds_manager, @@ -269,6 +273,7 @@ pub async fn handle_request( app: &AppHandle, app_data_root: &AppDataRoot, app_env_configs: &AppEnvConfigs, + artcraft_usage_tracker: &ArtcraftUsageTracker, provider_priority_store: &ProviderPriorityStore, task_database: &TaskDatabase, storyteller_creds_manager: &StorytellerCredentialManager, @@ -342,6 +347,13 @@ pub async fn handle_request( error!("Failed to create task in database: {:?}", err); // NB: Fail open, but find a way to flag this. } + + // TODO: We don't know that this is strictly the edit page. It could be the stage. It might also be nice to distinguish inpainting. + let num_images = request.image_count.map(|count| count as u16).unwrap_or(1); + if let Err(err) = artcraft_usage_tracker.record_image_generation(num_images, ArtcraftUsageType::ImageToResult, ArtcraftUsagePage::EditPage) { + // NB: Fail open. + warn!("Failed to report usage: {:?}", err); + } Ok(success_event) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/enqueue_image_inpaint_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/enqueue_image_inpaint_command.rs index b731ea208f..f132113aed 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/enqueue_image_inpaint_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/enqueue_image_inpaint_command.rs @@ -41,6 +41,8 @@ use storyteller_client::error::storyteller_error::StorytellerError; use storyteller_client::utils::api_host::ApiHost; use tauri::{AppHandle, Manager, State}; use tokens::tokens::media_files::MediaFileToken; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_tracker::ArtcraftUsageTracker; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_type::{ArtcraftUsagePage, ArtcraftUsageType}; #[derive(Deserialize, Debug, Copy, Clone)] #[serde(rename_all = "snake_case")] @@ -155,6 +157,7 @@ pub async fn enqueue_image_inpaint_command( request: EnqueueInpaintImageCommand, app_data_root: State<'_, AppDataRoot>, app_env_configs: State<'_, AppEnvConfigs>, + artcraft_usage_tracker: State<'_, ArtcraftUsageTracker>, provider_priority_store: State<'_, ProviderPriorityStore>, task_database: State<'_, TaskDatabase>, storyteller_creds_manager: State<'_, StorytellerCredentialManager>, @@ -170,6 +173,7 @@ pub async fn enqueue_image_inpaint_command( &app, &app_data_root, &app_env_configs, + &artcraft_usage_tracker, &provider_priority_store, &task_database, &storyteller_creds_manager, @@ -220,6 +224,7 @@ pub async fn handle_request( app: &AppHandle, app_data_root: &AppDataRoot, app_env_configs: &AppEnvConfigs, + artcraft_usage_tracker: &ArtcraftUsageTracker, provider_priority_store: &ProviderPriorityStore, task_database: &TaskDatabase, storyteller_creds_manager: &StorytellerCredentialManager, @@ -272,6 +277,13 @@ pub async fn handle_request( // NB: Fail open, but find a way to flag this. } + // TODO: We don't know that this is strictly the edit page. It could be the stage. It might also be nice to distinguish inpainting. + let num_images = request.image_count.map(|count| count as u16).unwrap_or(1); + if let Err(err) = artcraft_usage_tracker.record_image_generation(num_images, ArtcraftUsageType::ImageToResult, ArtcraftUsagePage::EditPage) { + // NB: Fail open. + warn!("Failed to report usage: {:?}", err); + } + Ok(success_event) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/enqueue_image_to_gaussian_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/enqueue_image_to_gaussian_command.rs index 5722069891..b58dd1b126 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/enqueue_image_to_gaussian_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/enqueue_image_to_gaussian_command.rs @@ -30,6 +30,8 @@ use serde::{Deserialize, Serialize}; use sqlite_tasks::queries::create_task::{create_task, CreateTaskArgs}; use tauri::{AppHandle, State}; use tokens::tokens::media_files::MediaFileToken; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_tracker::ArtcraftUsageTracker; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_type::{ArtcraftUsagePage, ArtcraftUsageType}; /// This is used in the Tauri command bridge. /// Don't change the serializations without coordinating with the frontend. @@ -97,6 +99,7 @@ pub async fn enqueue_image_to_gaussian_command( app: AppHandle, app_data_root: State<'_, AppDataRoot>, app_env_configs: State<'_, AppEnvConfigs>, + artcraft_usage_tracker: State<'_, ArtcraftUsageTracker>, task_database: State<'_, TaskDatabase>, storyteller_creds_manager: State<'_, StorytellerCredentialManager>, worldlabs_creds_manager: State<'_, WorldlabsCredentialManager>, @@ -108,6 +111,7 @@ pub async fn enqueue_image_to_gaussian_command( request, &app, &app_data_root, + &artcraft_usage_tracker, &task_database, &storyteller_creds_manager, &worldlabs_creds_manager, @@ -168,6 +172,7 @@ pub async fn handle_request( request: EnqueueImageToGaussianRequest, app: &AppHandle, app_data_root: &AppDataRoot, + artcraft_usage_tracker: &ArtcraftUsageTracker, task_database: &TaskDatabase, storyteller_creds_manager: &StorytellerCredentialManager, worldlabs_creds_manager: &WorldlabsCredentialManager, @@ -201,7 +206,12 @@ pub async fn handle_request( error!("Failed to create task in database: {:?}", err); // NB: Fail open, but find a way to flag this. } - + + if let Err(err) = artcraft_usage_tracker.record_object_generation(1, ArtcraftUsageType::ImageToResult, ArtcraftUsagePage::OtherPage) { + // NB: Fail open. + warn!("Failed to report usage: {:?}", err); + } + Ok(success_event) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_object/enqueue_image_to_3d_object_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_object/enqueue_image_to_3d_object_command.rs index e9cd6a0609..39bcf076d3 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_object/enqueue_image_to_3d_object_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_object/enqueue_image_to_3d_object_command.rs @@ -19,6 +19,8 @@ use log::{error, info, warn}; use serde_derive::{Deserialize, Serialize}; use tauri::{AppHandle, State}; use tokens::tokens::media_files::MediaFileToken; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_tracker::ArtcraftUsageTracker; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_type::{ArtcraftUsagePage, ArtcraftUsageType}; #[derive(Deserialize)] pub struct EnqueueImageTo3dObjectRequest { @@ -88,6 +90,7 @@ pub async fn enqueue_image_to_3d_object_command( request: EnqueueImageTo3dObjectRequest, app_env_configs: State<'_, AppEnvConfigs>, app_data_root: State<'_, AppDataRoot>, + artcraft_usage_tracker: State<'_, ArtcraftUsageTracker>, provider_priority_store: State<'_, ProviderPriorityStore>, task_database: State<'_, TaskDatabase>, storyteller_creds_manager: State<'_, StorytellerCredentialManager>, @@ -101,6 +104,7 @@ pub async fn enqueue_image_to_3d_object_command( &app, &app_env_configs, &app_data_root, + &artcraft_usage_tracker, &provider_priority_store, &task_database, &storyteller_creds_manager, @@ -171,6 +175,7 @@ pub async fn handle_request( app: &AppHandle, app_env_configs: &AppEnvConfigs, app_data_root: &AppDataRoot, + artcraft_usage_tracker: &ArtcraftUsageTracker, provider_priority_store: &ProviderPriorityStore, task_database: &TaskDatabase, storyteller_creds_manager: &StorytellerCredentialManager, @@ -203,6 +208,10 @@ pub async fn handle_request( // NB: Fail open, but find a way to flag this. } + if let Err(err) = artcraft_usage_tracker.record_object_generation(1, ArtcraftUsageType::ImageToResult, ArtcraftUsagePage::ObjectPage) { + // NB: Fail open. + warn!("Failed to report usage: {:?}", err); + } Ok(success_event) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_video/enqueue_image_to_video_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_video/enqueue_image_to_video_command.rs index a279cfbdc5..2c9af13843 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_video/enqueue_image_to_video_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_video/enqueue_image_to_video_command.rs @@ -14,6 +14,7 @@ use crate::core::events::generation_events::common::GenerationAction; use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent; use crate::core::events::generation_events::generation_enqueue_success_event::GenerationEnqueueSuccessEvent; use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_tracker::ArtcraftUsageTracker; use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::core::state::provider_priority::{Provider, ProviderPriorityStore}; use crate::core::state::task_database::TaskDatabase; @@ -27,6 +28,7 @@ use log::{error, info, warn}; use serde_derive::{Deserialize, Serialize}; use tauri::{AppHandle, State}; use tokens::tokens::media_files::MediaFileToken; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_type::{ArtcraftUsagePage, ArtcraftUsageType}; /// This is used in the Tauri command bridge. /// Don't change the serializations without coordinating with the frontend. @@ -175,6 +177,7 @@ pub async fn enqueue_image_to_video_command( app: AppHandle, app_env_configs: State<'_, AppEnvConfigs>, app_data_root: State<'_, AppDataRoot>, + artcraft_usage_tracker: State<'_, ArtcraftUsageTracker>, provider_priority_store: State<'_, ProviderPriorityStore>, task_database: State<'_, TaskDatabase>, grok_creds_manager: State<'_, GrokCredentialManager>, @@ -190,6 +193,7 @@ pub async fn enqueue_image_to_video_command( &app, &app_env_configs, &app_data_root, + &artcraft_usage_tracker, &provider_priority_store, &task_database, &grok_creds_manager, @@ -262,6 +266,7 @@ pub async fn handle_request( app: &AppHandle, app_env_configs: &AppEnvConfigs, app_data_root: &AppDataRoot, + artcraft_usage_tracker: &ArtcraftUsageTracker, provider_priority_store: &ProviderPriorityStore, task_database: &TaskDatabase, grok_creds_manager: &GrokCredentialManager, @@ -331,5 +336,10 @@ pub async fn handle_request( // NB: Fail open, but find a way to flag this. } + if let Err(err) = artcraft_usage_tracker.record_video_generation(1, ArtcraftUsageType::ImageToResult, ArtcraftUsagePage::VideoPage) { + // NB: Fail open. + warn!("Failed to report usage: {:?}", err); + } + Ok(success_event) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs index 48a578c3fc..7f19dc7ade 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs @@ -35,6 +35,8 @@ use sqlite_tasks::queries::create_task::{create_task, CreateTaskArgs}; use tauri::{AppHandle, State}; use tokens::tokens::media_files::MediaFileToken; use crate::core::api_adapters::aspect_ratio::common_aspect_ratio::CommonAspectRatio; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_tracker::ArtcraftUsageTracker; +use crate::core::state::artcraft_usage_tracker::artcraft_usage_type::{ArtcraftUsagePage, ArtcraftUsageType}; /// This is used in the Tauri command bridge. /// Don't change the serializations without coordinating with the frontend. @@ -179,6 +181,7 @@ pub async fn enqueue_text_to_image_command( app: AppHandle, app_data_root: State<'_, AppDataRoot>, app_env_configs: State<'_, AppEnvConfigs>, + artcraft_usage_tracker: State<'_, ArtcraftUsageTracker>, provider_priority_store: State<'_, ProviderPriorityStore>, task_database: State<'_, TaskDatabase>, mj_creds_manager: State<'_, MidjourneyCredentialManager>, @@ -197,6 +200,7 @@ pub async fn enqueue_text_to_image_command( request, &app, &app_data_root, + &artcraft_usage_tracker, &provider_priority_store, &task_database, &mj_creds_manager, @@ -267,6 +271,7 @@ pub async fn handle_request( request: EnqueueTextToImageRequest, app: &AppHandle, app_data_root: &AppDataRoot, + artcraft_usage_tracker: &ArtcraftUsageTracker, provider_priority_store: &ProviderPriorityStore, task_database: &TaskDatabase, mj_creds_manager: &MidjourneyCredentialManager, @@ -310,6 +315,24 @@ pub async fn handle_request( error!("Failed to create task in database: {:?}", err); // NB: Fail open, but find a way to flag this. } + + let num_images = request.number_images.map(|n| n as u16).unwrap_or(1); + + let is_image_to_image = request.image_media_tokens + .as_ref() + .map(|tokens| !tokens.is_empty()) + .unwrap_or(false); + + let usage_type = if is_image_to_image { + ArtcraftUsageType::ImageToResult + } else { + ArtcraftUsageType::TextToResult + }; + + if let Err(err) = artcraft_usage_tracker.record_image_generation(num_images, usage_type, ArtcraftUsagePage::ImagePage) { + // NB: Fail open. + warn!("Failed to report usage: {:?}", err); + } Ok(success_event) } diff --git a/crates/desktop/artcraft/src/core/state/artcraft_usage_tracker/artcraft_usage_tracker.rs b/crates/desktop/artcraft/src/core/state/artcraft_usage_tracker/artcraft_usage_tracker.rs index 4c510abe4d..fca429d97a 100644 --- a/crates/desktop/artcraft/src/core/state/artcraft_usage_tracker/artcraft_usage_tracker.rs +++ b/crates/desktop/artcraft/src/core/state/artcraft_usage_tracker/artcraft_usage_tracker.rs @@ -4,6 +4,8 @@ use log::{error, warn}; use crate::core::artcraft_error::ArtcraftError; use crate::core::state::artcraft_usage_tracker::artcraft_usage_type::{ArtcraftUsagePage, ArtcraftUsageType}; +/// Track generation volume so we know which areas to focus efforts on. +/// But also note: https://en.wikipedia.org/wiki/Survivorship_bias #[derive(Clone)] pub struct ArtcraftUsageTracker { data: Arc>, diff --git a/crates/desktop/artcraft/src/core/state/artcraft_usage_tracker/artcraft_usage_tracker_data.rs b/crates/desktop/artcraft/src/core/state/artcraft_usage_tracker/artcraft_usage_tracker_data.rs index 3b7f44c447..725c7609d9 100644 --- a/crates/desktop/artcraft/src/core/state/artcraft_usage_tracker/artcraft_usage_tracker_data.rs +++ b/crates/desktop/artcraft/src/core/state/artcraft_usage_tracker/artcraft_usage_tracker_data.rs @@ -1,4 +1,6 @@ +/// Track generation volume so we know which areas to focus efforts on. +/// But also note: https://en.wikipedia.org/wiki/Survivorship_bias #[derive(Clone)] pub struct ArtcraftUsageTrackerData { pub total_generation_count: u16, diff --git a/crates/desktop/artcraft/src/core/state/artcraft_usage_tracker/mod.rs b/crates/desktop/artcraft/src/core/state/artcraft_usage_tracker/mod.rs index 30be4a5d9c..8a61045b3d 100644 --- a/crates/desktop/artcraft/src/core/state/artcraft_usage_tracker/mod.rs +++ b/crates/desktop/artcraft/src/core/state/artcraft_usage_tracker/mod.rs @@ -1,3 +1,3 @@ pub mod artcraft_usage_tracker; -mod artcraft_usage_tracker_data; -mod artcraft_usage_type; \ No newline at end of file +pub mod artcraft_usage_type; +mod artcraft_usage_tracker_data; \ No newline at end of file