From a792ed9a093623d8d38976c3eaae8bc43b2753cf Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 16 May 2025 14:17:58 -0400 Subject: [PATCH] fal video wip, polling more complete --- .../fal/fal_background_removal_command.rs | 9 +- .../fal/fal_kling_image_to_video_command.rs | 246 ++++++++++++++++++ .../desktop/artcraft/src/commands/fal/mod.rs | 1 + crates/desktop/artcraft/src/lib.rs | 18 +- .../artcraft/src/state/fal/fal_task_queue.rs | 164 ++++++++++++ crates/desktop/artcraft/src/state/fal/mod.rs | 1 + .../src/threads/fal_task_polling_thread.rs | 187 +++++++++++++ crates/desktop/artcraft/src/threads/mod.rs | 1 + .../fal_client/src/fal_error_plus.rs | 3 + .../fal_client/src/model/enqueued_request.rs | 7 +- .../fal_client/src/model/fal_request_id.rs | 13 + .../api_clients/fal_client/src/model/mod.rs | 3 +- .../fal_client/src/requests/mod.rs | 4 +- .../api_clients/fal_client/src/utils/mod.rs | 3 +- .../src/utils/queue_status_checker.rs | 48 ++++ 15 files changed, 692 insertions(+), 16 deletions(-) create mode 100644 crates/desktop/artcraft/src/commands/fal/fal_kling_image_to_video_command.rs create mode 100644 crates/desktop/artcraft/src/state/fal/fal_task_queue.rs create mode 100644 crates/desktop/artcraft/src/threads/fal_task_polling_thread.rs create mode 100644 crates/lib/api_clients/fal_client/src/model/fal_request_id.rs create mode 100644 crates/lib/api_clients/fal_client/src/utils/queue_status_checker.rs diff --git a/crates/desktop/artcraft/src/commands/fal/fal_background_removal_command.rs b/crates/desktop/artcraft/src/commands/fal/fal_background_removal_command.rs index 076930642b..fffcf8d7b2 100644 --- a/crates/desktop/artcraft/src/commands/fal/fal_background_removal_command.rs +++ b/crates/desktop/artcraft/src/commands/fal/fal_background_removal_command.rs @@ -143,7 +143,7 @@ pub async fn fal_background_removal_command( //if let Err(err) = result { // error!("Failed to emit event: {:?}", err); //} - + Ok(FalBackgroundRemovalSuccessResponse { media_token: result.0.media_file.token, cdn_url: result.0.media_file.media_links.cdn_url, @@ -159,7 +159,6 @@ enum InnerError { AnyhowError(AnyhowError), StorytellerApiError(ApiError), DecodeError(DecodeError), - //ImageError(image::ImageError), IoError(std::io::Error), } @@ -187,12 +186,6 @@ impl From for InnerError { } } -//impl From for InnerError { -// fn from(value: image::ImageError) -> Self { -// Self::ImageError(value) -// } -//} -// impl From for InnerError { fn from(value: std::io::Error) -> Self { Self::IoError(value) diff --git a/crates/desktop/artcraft/src/commands/fal/fal_kling_image_to_video_command.rs b/crates/desktop/artcraft/src/commands/fal/fal_kling_image_to_video_command.rs new file mode 100644 index 0000000000..d5f2b3eb20 --- /dev/null +++ b/crates/desktop/artcraft/src/commands/fal/fal_kling_image_to_video_command.rs @@ -0,0 +1,246 @@ +use crate::commands::command_response_wrapper::{CommandErrorResponseWrapper, CommandErrorStatus, CommandResult, SerializeMarker}; +use crate::events::sendable_event_trait::SendableEvent; +use crate::events::sora::sora_image_enqueue_failure_event::SoraImageEnqueueFailureEvent; +use crate::events::sora::sora_image_enqueue_success_event::SoraImageEnqueueSuccessEvent; +use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::state::data_dir::trait_data_subdir::DataSubdir; +use crate::state::fal::fal_credential_manager::FalCredentialManager; +use crate::state::sora::read_sora_credentials_from_disk::read_sora_credentials_from_disk; +use crate::state::sora::sora_credential_holder::SoraCredentialHolder; +use crate::state::sora::sora_credential_manager::SoraCredentialManager; +use crate::state::sora::sora_task_queue::SoraTaskQueue; +use crate::state::storyteller::storyteller_credential_manager::StorytellerCredentialManager; +use crate::utils::get_url_file_extension::get_url_file_extension; +use crate::utils::simple_http_download::simple_http_download; +use crate::utils::simple_http_download_to_tempfile::simple_http_download_to_tempfile; +use anyhow::anyhow; +use base64::prelude::BASE64_STANDARD; +use base64::{DecodeError, Engine}; +use errors::{AnyhowError, AnyhowResult}; +use fal_client::fal_error_plus::FalErrorPlus; +use fal_client::requests::remove_background_rembg::remove_background_rembg; +use image::codecs::png::{CompressionType, FilterType, PngEncoder}; +use image::{DynamicImage, EncodableLayout, ImageReader}; +use log::{debug, error, info, warn}; +use mimetypes::mimetype_info::mimetype_info::MimetypeInfo; +use openai_sora_client::credentials::SoraCredentials; +use openai_sora_client::creds::credential_migration::CredentialMigrationRef; +use openai_sora_client::recipes::image_remix_with_session_auto_renew::{image_remix_with_session_auto_renew, ImageRemixAutoRenewRequest}; +use openai_sora_client::recipes::image_upload_from_file_with_session_auto_renew::{image_upload_from_file_with_session_auto_renew, ImageUploadFromFileAutoRenewRequest}; +use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session; +use openai_sora_client::requests::image_gen::common::{ImageSize, NumImages}; +use openai_sora_client::requests::image_gen::sora_image_gen_remix::{sora_image_gen_remix, SoraImageGenRemixRequest}; +use openai_sora_client::requests::upload::upload_media_from_bytes::sora_media_upload_from_bytes; +use openai_sora_client::requests::upload::upload_media_from_file::sora_media_upload_from_file; +use openai_sora_client::sora_error::SoraError; +use reqwest::Url; +use serde_derive::{Deserialize, Serialize}; +use std::fs::read_to_string; +use std::io::{Cursor, Read, Write}; +use std::ops::Add; +use std::path::PathBuf; +use std::time::Duration; +use storyteller_client::api_error::ApiError; +use storyteller_client::api_error::ApiError::InternalServerError; +use storyteller_client::media_files::get_media_file::{get_media_file, GetMediaFileSuccessResponse}; +use storyteller_client::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file; +use storyteller_client::utils::api_host::ApiHost; +use tauri::{AppHandle, Emitter, Manager, State}; +use tempfile::NamedTempFile; +use fal_client::requests::enqueue_kling_16_image_to_video::{enqueue_kling_16_image_to_video, Kling16Args, Kling16Duration}; +use filesys::file_read_bytes::file_read_bytes; +use tokens::tokens::media_files::MediaFileToken; + +#[derive(Deserialize)] +pub struct FalKlingImageToVideoRequest { + /// Image media file; the image to remove the background from. + pub image_media_token: Option, + + /// Base64-encoded image + pub base64_image: Option, +} + +#[derive(Serialize, Debug)] +#[serde(rename_all = "snake_case")] +pub enum SoraImageRemixErrorType { + /// Generic server error + ServerError, + /// No Fal API key available + NeedsFalApiKey, +} + + +#[tauri::command] +pub async fn fal_kling_image_to_video_command( + request: FalKlingImageToVideoRequest, + app_data_root: State<'_, AppDataRoot>, + fal_creds_manager: State<'_, FalCredentialManager>, + storyteller_creds_manager: State<'_, StorytellerCredentialManager>, +) -> CommandResult<(), SoraImageRemixErrorType, ()> { + + info!("fal_kling_image_to_video_command called; image media token: {:?}", request.image_media_token); + + let has_credentials = fal_creds_manager + .has_apparent_api_token() + .unwrap_or(true); // NB: Failures would be lock issues + + if !has_credentials { + warn!("No API key found"); + return Err(CommandErrorResponseWrapper { + status: CommandErrorStatus::Unauthorized, + error_message: Some("You need to set a FAL api key".to_string()), + error_type: Some(SoraImageRemixErrorType::NeedsFalApiKey), + error_details: None, + }); + } + + let result = image_to_video( + request, + &app_data_root, + &fal_creds_manager, + &storyteller_creds_manager + ).await; + + if let Err(err) = result { + error!("error: {:?}", err); + + //let event = SoraImageEnqueueFailureEvent {}; + //let result = event.send(&app); + //if let Err(err) = result { + // error!("Failed to emit event: {:?}", err); + //} + + let mut status = CommandErrorStatus::ServerError; + let mut error_type = SoraImageRemixErrorType::ServerError; + let mut error_message = "A server error occurred. Please try again. If it continues, please tell our staff about the problem."; + + //match (err) { + // _ => {}, + //} + + return Err(CommandErrorResponseWrapper { + status, + error_message: Some(error_message.to_string()), + error_type: Some(error_type), + error_details: None, + }) + } + + //let event = SoraImageEnqueueSuccessEvent {}; + //let result = event.send(&app); + //if let Err(err) = result { + // error!("Failed to emit event: {:?}", err); + //} + + Ok(().into()) +} + +#[derive(Debug)] +enum InnerError { + FalError(FalErrorPlus), + AnyhowError(AnyhowError), + StorytellerApiError(ApiError), + DecodeError(DecodeError), + IoError(std::io::Error), +} + +impl From for InnerError { + fn from(value: AnyhowError) -> Self { + Self::AnyhowError(value) + } +} + +impl From for InnerError { + fn from(value: FalErrorPlus) -> Self { + Self::FalError(value) + } +} + +impl From for InnerError { + fn from(value: ApiError) -> Self { + Self::StorytellerApiError(value) + } +} + +impl From for InnerError { + fn from(value: DecodeError) -> Self { + Self::DecodeError(value) + } +} + +impl From for InnerError { + fn from(value: std::io::Error) -> Self { + Self::IoError(value) + } +} + +pub async fn image_to_video( + request: FalKlingImageToVideoRequest, + app_data_root: &AppDataRoot, + fal_creds_manager: &FalCredentialManager, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> Result<(), InnerError> { + + let api_key = fal_creds_manager.get_key_required()?; + let creds = storyteller_creds_manager.get_credentials_required()?; + + let mut temp_download; + + if let Some(media_token) = request.image_media_token { + temp_download = download_media_file(&app_data_root, &media_token).await?; + + } else if let Some(base64_bytes) = request.base64_image { + temp_download = persist_base64(&app_data_root, base64_bytes).await?; + } else { + return Err(InnerError::AnyhowError(anyhow!("No image media token or base64 image provided"))); + } + + info!("Calling FAL image to video ..."); + + let filename = temp_download.path().to_path_buf(); + + let enqueued = enqueue_kling_16_image_to_video(Kling16Args { + image_path: filename, + api_key: &api_key, + duration: Kling16Duration::Default, + }).await?; + + // TODO: Handle event polling + + Ok(()) +} + +async fn download_media_file(app_data_root: &AppDataRoot, token: &MediaFileToken) -> Result { + let response = get_media_file(&ApiHost::Storyteller, token).await?; + + let media_file_url = &response.media_file.media_links.cdn_url; + + let extension_with_dot = get_url_file_extension(media_file_url) + .map(|ext| format!(".{}", ext)) + .unwrap_or_else(|| ".png".to_string()); + + //let filename = format!("{}{}", response.media_file.token.as_str(), extension_with_dot); + //let filename = app_data_root.downloads_dir().path().join(&filename); + + let mut file = app_data_root.temp_dir().new_named_temp_file_with_extension(&extension_with_dot)?; + + simple_http_download_to_tempfile(&media_file_url, &mut file).await?; + + Ok(file) // NB: Must return TempFile to not drop / delete it +} + +async fn persist_base64(app_data_root: &AppDataRoot, base64_image: String) -> Result { + let bytes = BASE64_STANDARD.decode(base64_image)?; + + let extension = MimetypeInfo::get_for_bytes(&bytes) + .map(|info| info.file_extension()) + .flatten() + .map(|ext| ext.extension_with_period().to_string()) + .unwrap_or_else(|| ".png".to_string()); + + let mut file = app_data_root.temp_dir().new_named_temp_file_with_extension(&extension)?; + + file.write_all(bytes.as_ref())?; + + Ok(file) // NB: Must return TempFile to not drop / delete it +} diff --git a/crates/desktop/artcraft/src/commands/fal/mod.rs b/crates/desktop/artcraft/src/commands/fal/mod.rs index c69b86c32b..d547c4ca54 100644 --- a/crates/desktop/artcraft/src/commands/fal/mod.rs +++ b/crates/desktop/artcraft/src/commands/fal/mod.rs @@ -1,3 +1,4 @@ pub mod fal_background_removal_command; +pub mod fal_kling_image_to_video_command; pub mod get_fal_api_key_command; pub mod set_fal_api_key_command; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index 82e491d46e..11964b5d80 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/crates/desktop/artcraft/src/lib.rs @@ -10,6 +10,7 @@ use tauri::Manager; use crate::commands::app_preferences::get_app_preferences_command::get_app_preferences_command; use crate::commands::app_preferences::update_app_preference_command::update_app_preferences_command; use crate::commands::fal::fal_background_removal_command::fal_background_removal_command; +use crate::commands::fal::fal_kling_image_to_video_command::fal_kling_image_to_video_command; use crate::commands::fal::get_fal_api_key_command::get_fal_api_key_command; use crate::commands::fal::set_fal_api_key_command::set_fal_api_key_command; use crate::commands::flip_image::flip_image; @@ -22,12 +23,14 @@ use crate::commands::sora::sora_logout_command::sora_logout_command; use crate::state::app_preferences::app_preferences_manager::load_app_preferences_or_default; use crate::state::data_dir::app_data_root::AppDataRoot; use crate::state::fal::fal_credential_manager::FalCredentialManager; +use crate::state::fal::fal_task_queue::FalTaskQueue; use crate::state::main_window_position::MainWindowPosition; use crate::state::main_window_size::MainWindowSize; use crate::state::sora::sora_credential_manager::SoraCredentialManager; use crate::state::sora::sora_task_queue::SoraTaskQueue; use crate::state::storyteller::storyteller_credential_manager::StorytellerCredentialManager; use crate::threads::discord_presence_thread::discord_presence_thread; +use crate::threads::fal_task_polling_thread::fal_task_polling_thread; use crate::threads::main_window_thread::main_window_thread::main_window_thread; use crate::threads::sora_task_polling_thread::sora_task_polling_thread; use crate::utils::webview_unsafe::webview_unsafe_for_app; @@ -59,11 +62,15 @@ pub fn run() { println!("Attempting to read existing fal credentials..."); let fal_creds_manager = FalCredentialManager::initialize_from_disk_infallible(&app_data_root); + let fal_creds_manager_2 = fal_creds_manager.clone(); // Other state let sora_task_queue = SoraTaskQueue::new(); let sora_task_queue_2 = sora_task_queue.clone(); + let fal_task_queue = FalTaskQueue::new(); + let fal_task_queue_2 = fal_task_queue.clone(); + println!("Initializing backend runtime..."); tauri::Builder::default() @@ -97,13 +104,19 @@ pub fn run() { eprintln!("Error setting webview unsafe: {:?}", err); } + // TODO(bt): Clean this up. We can just clone at the callsite. Also clean initialization let app_2 = app.clone(); let app_3 = app.clone(); let app_4 = app.clone(); + let app_5 = app.clone(); let app_data_root_3 = app_data_root_2.clone(); let app_data_root_4 = app_data_root_2.clone(); + let app_data_root_5 = app_data_root_2.clone(); let sora_creds_manager_3 = sora_creds_manager_2.clone(); let storyteller_creds_manager_3 = storyteller_creds_manager.clone(); + let storyteller_creds_manager_4 = storyteller_creds_manager.clone(); + let fal_creds_manager_3 = fal_creds_manager_2.clone(); + let fal_task_queue_3 = fal_task_queue_2.clone(); match MainWindowSize::from_filesystem_configs(&app_data_root_3) { Ok(None) => {} @@ -136,18 +149,21 @@ pub fn run() { tauri::async_runtime::spawn(main_window_thread(app_3, app_data_root_3, storyteller_creds_manager_2)); tauri::async_runtime::spawn(sora_task_polling_thread(app_4, app_data_root_4, sora_creds_manager_3, storyteller_creds_manager_3, sora_task_queue_2)); tauri::async_runtime::spawn(discord_presence_thread()); + tauri::async_runtime::spawn(fal_task_polling_thread(app_5, app_data_root_5, fal_creds_manager_3, storyteller_creds_manager_4, fal_task_queue_3)); Ok(()) }) .manage(app_data_root) .manage(app_preferences) - .manage(sora_creds_manager) .manage(fal_creds_manager) + .manage(fal_task_queue) + .manage(sora_creds_manager) .manage(sora_task_queue) .manage(storyteller_creds_manager_3) .invoke_handler(tauri::generate_handler![ check_sora_session_command, fal_background_removal_command, + fal_kling_image_to_video_command, flip_image, get_app_preferences_command, get_fal_api_key_command, diff --git a/crates/desktop/artcraft/src/state/fal/fal_task_queue.rs b/crates/desktop/artcraft/src/state/fal/fal_task_queue.rs new file mode 100644 index 0000000000..e733c886e1 --- /dev/null +++ b/crates/desktop/artcraft/src/state/fal/fal_task_queue.rs @@ -0,0 +1,164 @@ +use chrono::{DateTime, Utc}; +use errors::AnyhowResult; +use fal_client::model::enqueued_request::EnqueuedRequest; +use fal_client::model::fal_request_id::FalRequestId; +use indexmap::IndexMap; +use std::collections::HashSet; +use std::sync::{Arc, RwLock}; + +#[derive(Clone)] +pub struct FalTaskDetails { + pub enqueue_time: DateTime, + pub enqueued_request: EnqueuedRequest, +} + +impl FalTaskDetails { + pub fn for_request(request: EnqueuedRequest) -> Self { + Self { + enqueue_time: Utc::now(), + enqueued_request: request, + } + } +} + +#[derive(Clone)] +pub struct FalTaskQueue { + // Insertion-order map of task IDs to task details + queue: Arc>> +} + +impl FalTaskQueue { + pub fn new() -> Self { + Self { + queue: Arc::new(RwLock::new(IndexMap::new())) + } + } + + /// Determine if queue is empty + /// This is an O(1) operation. + pub fn is_empty(&self) -> AnyhowResult { + match self.queue.read() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)) + } + Ok(lock) => Ok(lock.is_empty()) + } + } + + /// Get length of set + /// This is an O(1) operation. + pub fn len(&self) -> AnyhowResult { + match self.queue.read() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)) + } + Ok(lock) => Ok(lock.len()) + } + } + + /// Get the task at an index. + /// This is an O(1) operation. + pub fn get_index(&self, i: usize) -> AnyhowResult> { + match self.queue.read() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)) + } + Ok(lock) => { + if i >= lock.len() { + Ok(None) + } else { + match lock.get_index(i) { + None => Ok(None), + Some((_key, value)) => { + Ok(Some(value.clone())) + } + } + } + } + } + } + + /// Determine if the key is in the queue + /// This is an O(1) operation. + pub fn contains_key(&self, request_id: &FalRequestId) -> AnyhowResult { + match self.queue.read() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)) + } + Ok(lock) => Ok(lock.contains_key(request_id)) + } + } + + /// Get the first enqueued task + /// This is an O(1) operation. + pub fn first(&self) -> AnyhowResult> { + match self.queue.read() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)) + } + Ok(lock) => match lock.first() { + None => Ok(None), + Some((request_id, task_details)) => Ok(Some(task_details.clone())) + } + } + } + + /// Get the last enqueued task + /// This is an O(1) operation. + pub fn last(&self) -> AnyhowResult> { + match self.queue.read() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)) + } + Ok(lock) => match lock.last() { + None => Ok(None), + Some((request_id, task_details)) => Ok(Some(task_details.clone())) + } + } + } + + /// Insert a new task. If a task already exists, we update its value. + /// This is an O(1) operation. + pub fn insert(&self, request: &EnqueuedRequest) -> AnyhowResult> { + match self.queue.write() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)) + } + Ok(mut lock) => { + // NB: Probably fine to compute TaskDetails here for now. We may eventually expose this if we include additional info. + let task_details = FalTaskDetails::for_request(request.clone()); + Ok(lock.insert(request.request_id.clone(), task_details)) + } + } + } + + /// Remove a task by ID. + /// This is an O(n) operation. + pub fn remove(&self, request_id: &FalRequestId) -> AnyhowResult> { + match self.queue.write() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)) + } + Ok(mut lock) => { + Ok(lock.shift_remove(request_id)) + } + } + } + + /// Remove a list of tasks by ID. + /// This is an O(n) operation. + pub fn remove_list(&self, request_ids: &[&FalRequestId]) -> AnyhowResult<()> { + match self.queue.write() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)) + } + Ok(mut lock) => { + // TODO: Cleaner way to do this? + let mut set : HashSet<&FalRequestId, _> = HashSet::with_capacity(request_ids.len()); + set.extend(request_ids); + lock.retain(|k, v| !set.contains(k)); + Ok(()) + } + } + } +} diff --git a/crates/desktop/artcraft/src/state/fal/mod.rs b/crates/desktop/artcraft/src/state/fal/mod.rs index fd55133c18..f9f5484671 100644 --- a/crates/desktop/artcraft/src/state/fal/mod.rs +++ b/crates/desktop/artcraft/src/state/fal/mod.rs @@ -1,2 +1,3 @@ pub mod fal_credential_manager; +pub mod fal_task_queue; pub mod read_fal_api_key_from_disk; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/threads/fal_task_polling_thread.rs b/crates/desktop/artcraft/src/threads/fal_task_polling_thread.rs new file mode 100644 index 0000000000..150e0fd290 --- /dev/null +++ b/crates/desktop/artcraft/src/threads/fal_task_polling_thread.rs @@ -0,0 +1,187 @@ +use crate::events::sendable_event_trait::SendableEvent; +use crate::events::sora::sora_image_generation_complete_event::SoraImageGenerationCompleteEvent; +use crate::events::sora::sora_image_generation_failed_event::SoraImageGenerationFailedEvent; +use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::state::data_dir::trait_data_subdir::DataSubdir; +use crate::state::fal::fal_credential_manager::FalCredentialManager; +use crate::state::fal::fal_task_queue::FalTaskQueue; +use crate::state::sora::sora_credential_manager::SoraCredentialManager; +use crate::state::storyteller::storyteller_credential_manager::StorytellerCredentialManager; +use errors::AnyhowResult; +use log::{error, info, warn}; +use openai_sora_client::requests::image_gen::image_gen_status::{Generation, TaskId, TaskStatus}; +use reqwest::Url; +use serde_derive::Serialize; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; +use storyteller_client::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file; +use storyteller_client::utils::api_host::ApiHost; +use tauri::{AppHandle, Emitter}; +use tempdir::TempDir; +use fal_client::fal_error_plus::FalErrorPlus; +use fal_client::utils::queue_status_checker::QueueStatusChecker; +use tokens::tokens::media_files::MediaFileToken; + +pub async fn fal_task_polling_thread( + app_handle: AppHandle, + app_data_root: AppDataRoot, + fal_creds_manager: FalCredentialManager, + storyteller_creds_manager: StorytellerCredentialManager, + fal_task_queue: FalTaskQueue, +) -> ! { + loop { + let res = polling_loop( + &app_handle, + &fal_creds_manager, + &storyteller_creds_manager, + &fal_task_queue, + &app_data_root, + ).await; + if let Err(err) = res { + error!("An error occurred: {:?}", err); + } + tokio::time::sleep(std::time::Duration::from_millis(30_000)).await; + } +} + +async fn polling_loop( + app_handle: &AppHandle, + fal_creds_manager: &FalCredentialManager, + storyteller_creds_manager: &StorytellerCredentialManager, + fal_task_queue: &FalTaskQueue, + app_data_root: &AppDataRoot, +) -> AnyhowResult<()> { + loop { + if fal_task_queue.is_empty()? { + // No need to poll if we don't have pending tasks. + tokio::time::sleep(std::time::Duration::from_millis(1_000)).await; + continue; + } + + info!("Task queue has {} pending tasks.", fal_task_queue.len()?); + + let key = fal_creds_manager.get_key_required()?; + + let queue_status_checker = QueueStatusChecker::new(key); + + // TODO: Clear completed tasks. + // let mut succeeded_tasks = Vec::new(); + // let mut failed_tasks = Vec::new(); + + let len = fal_task_queue.len()?; + + for i in 0..len { + let task = match fal_task_queue.get_index(i)? { + Some(task) => task, + None => { + warn!("Failed to get task ID at index {}", i); + break; + } + }; + + info!("Checking task status for task ID: {:?}", &task.enqueued_request.request_id); + + let result = queue_status_checker.check_status(&task.enqueued_request).await; + + match result { + Err(err) => { + warn!("Error checking job status: {:?}", err); + tokio::time::sleep(std::time::Duration::from_millis(2_000)).await; + continue; + } + Ok(status) => { + info!("Status: {:?}", status); + } + } + } + + + +// // TODO: The cursoring logic likely needs to improve. +// for task in response.task_responses { +// let status = TaskStatus::from_str(&task.status); +// +// match status { +// TaskStatus::Succeeded => { +// succeeded_tasks.push(task); +// } +// TaskStatus::Failed => { +// failed_tasks.push(task); +// } +// TaskStatus::Queued => {} +// TaskStatus::Running => {} +// TaskStatus::Unknown(_) => {} +// } +// } +// +// for task in failed_tasks.iter() { +// if fal_task_queue.contains_key(&task.id)? { +// let event = SoraImageGenerationFailedEvent { +// prompt: task.prompt.clone(), +// }; +// event.send(&app_handle)?; +// } +// } +// +// let failed_task_ids: Vec<&TaskId> = failed_tasks +// .iter() +// .map(|task| &task.id) +// .collect(); +// +// fal_task_queue.remove_list(&failed_task_ids)?; +// +// let creds = storyteller_creds_manager.get_credentials_required()?; +// let api_host = ApiHost::Storyteller; +// +// for task in succeeded_tasks.iter() { +// if !fal_task_queue.contains_key(&task.id)? { +// continue; +// } +// info!("Task succeeded: {:?}", task.id); +// +// for (i, generation) in task.generations.iter().enumerate() { +// info!("Downloading generated file..."); +// let download_path = download_generation(generation, &app_data_root).await?; +// +// info!("Uploading to backend..."); +// let result = upload_image_media_file_from_file(&api_host, Some(&creds), download_path).await?; +// +// info!("Uploaded to API backend: {:?}", result.media_file_token); +// +// let event = SoraImageGenerationCompleteEvent { +// media_file_token: result.media_file_token, +// }; +// +// event.send(&app_handle)?; +// } +// } +// +// let succeeded_task_ids : Vec<&TaskId> = succeeded_tasks +// .iter() +// .map(|task| &task.id) +// .collect(); +// +// fal_task_queue.remove_list(&succeeded_task_ids)?; + + tokio::time::sleep(std::time::Duration::from_millis(6_000)).await; + } +} + +async fn download_generation(generation: &Generation, app_data_root: &AppDataRoot) -> AnyhowResult { + let url = Url::parse(&generation.url)?; + + let response = reqwest::get(&generation.url).await?; + let image_bytes = response.bytes().await?; + + let ext = url.path().split(".").last().unwrap_or("png"); + + let tempdir = app_data_root.temp_dir().path(); + let download_filename = format!("{}.{}", generation.id, 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/threads/mod.rs b/crates/desktop/artcraft/src/threads/mod.rs index eaae9fdf2e..5b72269773 100644 --- a/crates/desktop/artcraft/src/threads/mod.rs +++ b/crates/desktop/artcraft/src/threads/mod.rs @@ -1,3 +1,4 @@ pub mod discord_presence_thread; +pub mod fal_task_polling_thread; pub mod main_window_thread; pub mod sora_task_polling_thread; \ No newline at end of file diff --git a/crates/lib/api_clients/fal_client/src/fal_error_plus.rs b/crates/lib/api_clients/fal_client/src/fal_error_plus.rs index 5168884a73..57259243b6 100644 --- a/crates/lib/api_clients/fal_client/src/fal_error_plus.rs +++ b/crates/lib/api_clients/fal_client/src/fal_error_plus.rs @@ -10,6 +10,8 @@ pub enum FalErrorPlus { AnyhowError(anyhow::Error), /// URL parse errors. UrlParseError(url::ParseError), + /// An endpoint we don't support yet. + UnhandledEndpoint(String), } impl Display for FalErrorPlus { @@ -18,6 +20,7 @@ impl Display for FalErrorPlus { Self::FalError(err) => write!(f, "FalErrorPlus::FalError: {:?}", err), Self::AnyhowError(err) => write!(f, "FalErrorPlus::AnyhowError: {:?}", err), Self::UrlParseError(err) => write!(f, "FalErrorPlus::UrlParseError: {:?}", err), + Self::UnhandledEndpoint(endpoint) => write!(f, "FalErrorPlus::UnhandledEndpoint: {:?}", endpoint), } } } diff --git a/crates/lib/api_clients/fal_client/src/model/enqueued_request.rs b/crates/lib/api_clients/fal_client/src/model/enqueued_request.rs index 4629349867..4fce8291ea 100644 --- a/crates/lib/api_clients/fal_client/src/model/enqueued_request.rs +++ b/crates/lib/api_clients/fal_client/src/model/enqueued_request.rs @@ -1,5 +1,6 @@ use crate::creds::fal_api_key::FalApiKey; use crate::model::fal_endpoint::FalEndpoint; +use crate::model::fal_request_id::FalRequestId; use errors::AnyhowResult; use fal::queue::{Queue, QueueResponse}; use reqwest::Client; @@ -8,7 +9,7 @@ use serde::de::DeserializeOwned; #[derive(Clone, Debug)] pub struct EnqueuedRequest { pub fal_endpoint: FalEndpoint, - pub request_id: String, + pub request_id: FalRequestId, pub response_url: String, pub status_url: String, pub cancel_url: String, @@ -19,7 +20,7 @@ impl EnqueuedRequest { let endpoint = FalEndpoint::from_queue_response(queue); Ok(Self { fal_endpoint: endpoint.clone(), - request_id: queue.payload.request_id.clone(), + request_id: FalRequestId(queue.payload.request_id.clone()), response_url: queue.payload.response_url.clone(), status_url: queue.payload.status_url.clone(), cancel_url: queue.payload.cancel_url.clone(), @@ -30,7 +31,7 @@ impl EnqueuedRequest { let client = Client::new(); let endpoint = self.fal_endpoint.url(); let payload = QueueResponse { - request_id: self.request_id.clone(), + request_id: self.request_id.0.clone(), response_url: self.response_url.clone(), status_url: self.status_url.clone(), cancel_url: self.cancel_url.clone(), diff --git a/crates/lib/api_clients/fal_client/src/model/fal_request_id.rs b/crates/lib/api_clients/fal_client/src/model/fal_request_id.rs new file mode 100644 index 0000000000..254cf8171f --- /dev/null +++ b/crates/lib/api_clients/fal_client/src/model/fal_request_id.rs @@ -0,0 +1,13 @@ +use std::fmt::Display; +use serde::{Deserialize, Serialize}; +use crate::model::fal_endpoint::FalEndpoint; + +/// A strongly typed request ID for the Fal API. +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Serialize, Deserialize)] +pub struct FalRequestId(pub String); + +impl Display for FalRequestId { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + Display::fmt(&self.0, f) + } +} diff --git a/crates/lib/api_clients/fal_client/src/model/mod.rs b/crates/lib/api_clients/fal_client/src/model/mod.rs index d23e0557fe..07ae4fa566 100644 --- a/crates/lib/api_clients/fal_client/src/model/mod.rs +++ b/crates/lib/api_clients/fal_client/src/model/mod.rs @@ -1,2 +1,3 @@ pub mod enqueued_request; -pub mod fal_endpoint; \ No newline at end of file +pub mod fal_endpoint; +pub mod fal_request_id; \ No newline at end of file diff --git a/crates/lib/api_clients/fal_client/src/requests/mod.rs b/crates/lib/api_clients/fal_client/src/requests/mod.rs index df837499b0..5490aa7586 100644 --- a/crates/lib/api_clients/fal_client/src/requests/mod.rs +++ b/crates/lib/api_clients/fal_client/src/requests/mod.rs @@ -1,2 +1,2 @@ -pub mod remove_background_rembg; -mod enqueue_kling_16_image_to_video; \ No newline at end of file +pub mod enqueue_kling_16_image_to_video; +pub mod remove_background_rembg; \ No newline at end of file diff --git a/crates/lib/api_clients/fal_client/src/utils/mod.rs b/crates/lib/api_clients/fal_client/src/utils/mod.rs index 18c0af1eb4..31c3f2080e 100644 --- a/crates/lib/api_clients/fal_client/src/utils/mod.rs +++ b/crates/lib/api_clients/fal_client/src/utils/mod.rs @@ -1,2 +1,3 @@ pub mod clone_queue_payload; -pub mod clone_queue_response; \ No newline at end of file +pub mod clone_queue_response; +pub mod queue_status_checker; \ No newline at end of file diff --git a/crates/lib/api_clients/fal_client/src/utils/queue_status_checker.rs b/crates/lib/api_clients/fal_client/src/utils/queue_status_checker.rs new file mode 100644 index 0000000000..256a16828a --- /dev/null +++ b/crates/lib/api_clients/fal_client/src/utils/queue_status_checker.rs @@ -0,0 +1,48 @@ +use fal::endpoints::fal_ai::kling_video::v1_6::pro::effects::I2VOutput; +use fal::FalError; +use fal::prelude::QueueResponse; +use fal::queue::{Queue, QueueStatus}; +use reqwest::Client; +use crate::creds::fal_api_key::FalApiKey; +use crate::fal_error_plus::FalErrorPlus; +use crate::model::enqueued_request::EnqueuedRequest; +use crate::model::fal_endpoint::FalEndpoint; + +pub struct QueueStatusChecker { + api_key: FalApiKey, +} + +impl QueueStatusChecker { + pub fn new(api_key: FalApiKey) -> Self { + Self { api_key } + } + + pub async fn check_status(&self, request: &EnqueuedRequest) -> Result { + let client = Client::new(); + let endpoint = request.fal_endpoint.url(); + + let payload = QueueResponse { + request_id: request.request_id.0.clone(), + response_url: request.response_url.clone(), + status_url: request.status_url.clone(), + cancel_url: request.cancel_url.clone(), + }; + + let queue ; + + match &request.fal_endpoint { + FalEndpoint::Kling16ImageToVideo => { + // TODO: Cleanup / make concise + let f : Queue = Queue::new(client, endpoint, self.api_key.0.clone(), payload); + queue = f; + } + FalEndpoint::Other(endpoint) => { + return Err(FalErrorPlus::UnhandledEndpoint(format!("Unsupported endpoint: {}", endpoint))); + } + } + + let status = queue.status(false).await?; + + Ok(status) + } +} \ No newline at end of file