diff --git a/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs b/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs index e6b82f2d2e..fcaef9e389 100644 --- a/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs +++ b/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs @@ -78,8 +78,10 @@ pub async fn handle_tauri_startup( tauri::async_runtime::spawn(midjourney_long_polling_thread( app.clone(), app_env_configs.clone(), + root.clone(), task_database.clone(), mj_creds_manager.clone(), + storyteller_creds_manager.clone(), )); spawn_discord_presence_thread()?; diff --git a/crates/desktop/artcraft/src/services/midjourney/mod.rs b/crates/desktop/artcraft/src/services/midjourney/mod.rs index ad664c3b7d..9e15c7b3a9 100644 --- a/crates/desktop/artcraft/src/services/midjourney/mod.rs +++ b/crates/desktop/artcraft/src/services/midjourney/mod.rs @@ -1,4 +1,5 @@ -pub mod threads; -pub mod state; pub mod commands; +pub mod state; +pub mod threads; +pub mod utils; pub mod windows; \ No newline at end of file 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 3d4d094a90..9ec26a5e4a 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 @@ -1,31 +1,47 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::data_dir::trait_data_subdir::DataSubdir; use crate::core::state::task_database::TaskDatabase; use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; -use errors::AnyhowResult; -use log::error; -use midjourney_client::client::midjourney_hostname::MidjourneyHostname; -use midjourney_client::endpoints::imagine::{imagine, ImagineRequest}; -use tauri::AppHandle; use enums::common::generation_provider::GenerationProvider; use enums::tauri::tasks::task_status::TaskStatus; +use errors::AnyhowResult; +use log::{error, info}; +use midjourney_client::client::midjourney_hostname::MidjourneyHostname; +use midjourney_client::endpoints::imagine::{imagine, ImagineRequest}; use midjourney_client::utils::get_image_url::get_image_url; -use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs}; +use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task}; +use std::collections::HashMap; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; +use tauri::AppHandle; +use url::Url; +use storyteller_client::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs}; +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::services::midjourney::utils::download_midjourney_image::download_midjourney_image; /// This thread is responsible for picking up tasks that fell through the cracks of /// the faster websocket thread. pub async fn midjourney_long_polling_thread( app_handle: AppHandle, app_env_configs: AppEnvConfigs, + app_data_root: AppDataRoot, task_database: TaskDatabase, creds: MidjourneyCredentialManager, + storyteller_creds_manager: StorytellerCredentialManager, ) -> ! { loop { let res = polling_loop( &app_handle, &app_env_configs, + &app_data_root, &task_database, &creds, + &storyteller_creds_manager, ).await; if let Err(err) = res { error!("An error occurred: {:?}", err); @@ -38,8 +54,10 @@ pub async fn midjourney_long_polling_thread( async fn polling_loop( app_handle: &AppHandle, app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, task_database: &TaskDatabase, creds: &MidjourneyCredentialManager, + storyteller_creds_manager: &StorytellerCredentialManager, ) -> AnyhowResult<()> { loop { if !creds.session_appears_active()? { @@ -47,6 +65,9 @@ async fn polling_loop( continue; } + // TODO: Graceful wait, fix this long function body + let storyteller_creds = storyteller_creds_manager.get_credentials_required()?; + let cookies = creds.maybe_copy_cookie_store()?; let cookies = match cookies { @@ -80,40 +101,108 @@ async fn polling_loop( page_size: None, }).await?; - let job_ids = result.items.iter() - .filter_map(|item| item.id.as_ref()) - .map(|id| id.to_string()) - .collect::>(); + let midjourney_items = result.items; - let tasks = list_tasks_by_provider_and_tokens(ListTasksArgs { + let midjourney_items_by_id = { + let mut hash = HashMap::new(); + for item in midjourney_items.iter() { + if let Some(id) = &item.id { + hash.insert(id.to_string(), item.clone()); + } + } + hash + }; + + //let midjourney_item_ids = midjourney_items.iter() + // .filter_map(|item| item.id.as_ref()) + // .map(|id| id.to_string()) + // .collect::>(); + + let midjourney_item_ids = midjourney_items_by_id + .keys() + .map(|id| id.to_string()) + .collect(); + + let local_tasks = list_tasks_by_provider_and_tokens(ListTasksArgs { db: task_database.get_connection(), provider: GenerationProvider::Midjourney, - provider_job_ids: Some(job_ids), + provider_job_ids: Some(midjourney_item_ids), }).await?; - let tasks = tasks.tasks; + let local_tasks = local_tasks.tasks; - for task in tasks { - let job_id = match task.provider_job_id { - Some(job_id) => job_id, + // Map of Midjourney Job ID to Local Task. + let local_tasks_by_midjourney_job_id = local_tasks.iter() + .filter(|task| match task.status { + TaskStatus::Pending => true, + TaskStatus::Started => true, + TaskStatus::AttemptFailed => true, + _ => false, + }) + .filter_map(|task| { + if let Some(provider_job_id) = &task.provider_job_id { + Some((provider_job_id.clone(), task.clone())) + } else { + None + } + }) + .collect::>(); + + // TODO: If we introduce another job polling mechanism, we may need to handle concurrency. + + for (midjourney_job_id, local_task) in local_tasks_by_midjourney_job_id.iter() { + let midjourney_item = match midjourney_items_by_id.get(midjourney_job_id) { + Some(item) => item, None => continue, }; - match task.status { - TaskStatus::Pending => {}, // Fall-through - TaskStatus::Started => {}, // Fall-through - TaskStatus::AttemptFailed => {}, // Fall-through - _ => continue, - } - for index in 0..4 { - let url = get_image_url(&job_id, index)?; + info!("Downloading generated Midjourney file..."); + let download_path = download_midjourney_image(midjourney_job_id, index, app_data_root).await?; - println!("Image URL for job {} index {}: {}", job_id, index, url); + info!("Uploading to backend..."); + + let result = upload_image_media_file_from_file(UploadImageFromFileArgs { + api_host: &app_env_configs.storyteller_host, + maybe_creds: Some(&storyteller_creds), + path: download_path, + is_intermediate_system_file: false, + }).await?; + + info!("Uploaded to API backend: {:?}", result.media_file_token); + + tokio::time::sleep(std::time::Duration::from_millis(100)).await; } + + let event = GenerationCompleteEvent { + //media_file_token: result.media_file_token, + action: Some(GenerationAction::GenerateImage), + service: GenerationServiceProvider::Midjourney, + model: None, + }; + + if let Err(err) = event.send(&app_handle) { + error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open + } + + tokio::time::sleep(std::time::Duration::from_millis(2_000)).await; } - //println!("Response: {:?}\n\n", tasks); + //for task in tasks { + // let job_id = match task.provider_job_id { + // Some(job_id) => job_id, + // None => continue, + // }; + // match task.status { + // TaskStatus::Pending => {}, // Fall-through + // TaskStatus::Started => {}, // Fall-through + // TaskStatus::AttemptFailed => {}, // Fall-through + // _ => { + // tasks_by_provider_job_id.remove(&job_id); + // continue + // }, + // } + //} tokio::time::sleep(std::time::Duration::from_millis(60_000)).await; } diff --git a/crates/desktop/artcraft/src/services/midjourney/utils/download_midjourney_image.rs b/crates/desktop/artcraft/src/services/midjourney/utils/download_midjourney_image.rs new file mode 100644 index 0000000000..6bcebd851e --- /dev/null +++ b/crates/desktop/artcraft/src/services/midjourney/utils/download_midjourney_image.rs @@ -0,0 +1,27 @@ +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::data_dir::trait_data_subdir::DataSubdir; +use errors::AnyhowResult; +use midjourney_client::utils::get_image_url::get_image_url; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; +use url::Url; + +pub async fn download_midjourney_image(job_id: &str, index: u8, app_data_root: &AppDataRoot) -> AnyhowResult { + let url = get_image_url(job_id, index)?; + let parsed_url = Url::parse(&url)?; + + let response = reqwest::get(&url).await?; + let image_bytes = response.bytes().await?; + + let ext = parsed_url.path().split(".").last().unwrap_or("png"); + + let tempdir = app_data_root.temp_dir().path(); + let download_filename = format!("{}_{}.{}", job_id, index, ext); + let download_path = tempdir.join(download_filename); + + let mut file = File::create(&download_path)?; + file.write_all(&image_bytes)?; + + Ok(download_path) +} diff --git a/crates/desktop/artcraft/src/services/midjourney/utils/mod.rs b/crates/desktop/artcraft/src/services/midjourney/utils/mod.rs new file mode 100644 index 0000000000..acf0cc90dc --- /dev/null +++ b/crates/desktop/artcraft/src/services/midjourney/utils/mod.rs @@ -0,0 +1 @@ +pub mod download_midjourney_image; \ No newline at end of file