diff --git a/Cargo.lock b/Cargo.lock index 08a8b84a6c..d58d36ed77 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4254,6 +4254,7 @@ dependencies = [ "futures 0.3.31", "testing", "tokio", + "url", "workspace-build-acceleration", ] diff --git a/crates/desktop/artcraft/src/commands/command_response_wrapper.rs b/crates/desktop/artcraft/src/commands/command_response_wrapper.rs index 610f129efb..0999edfc21 100644 --- a/crates/desktop/artcraft/src/commands/command_response_wrapper.rs +++ b/crates/desktop/artcraft/src/commands/command_response_wrapper.rs @@ -73,7 +73,7 @@ impl CommandSuccessResponseWrapper<()> { /// `Serialize`. Rust's `"specialization"`, which would let us account for this, is an unstable, nightly-only /// feature. So instead we manually mark all types other than `()` with `SerializeMarker` as a hack to /// distinguish the two. -trait SerializeMarker : Serialize {} +pub trait SerializeMarker : Serialize {} // NB: () is Serialize! So we can't broadly implement just `T: Serialize`. impl From<()> for CommandSuccessResponseWrapper<()> { 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 new file mode 100644 index 0000000000..f0f96a36e8 --- /dev/null +++ b/crates/desktop/artcraft/src/commands/fal/fal_background_removal_command.rs @@ -0,0 +1,215 @@ +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 base64::prelude::BASE64_STANDARD; +use base64::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, ImageReader}; +use log::{debug, error, info, warn}; +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; +use std::ops::Add; +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 tokens::tokens::media_files::MediaFileToken; + +#[derive(Deserialize)] +pub struct FalBackgroundRemovalRequest { + /// Image media file; the image to remove the background from. + pub image_media_token: MediaFileToken, +} + +#[derive(Serialize)] +pub struct FalBackgroundRemovalSuccessResponse { + /// Result media token + pub media_token: MediaFileToken, + /// Result URL + pub cdn_url: Url, +} + +impl SerializeMarker for FalBackgroundRemovalSuccessResponse {} + +#[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_background_removal_command( + app: AppHandle, + request: FalBackgroundRemovalRequest, + app_data_root: State<'_, AppDataRoot>, + fal_creds_manager: State<'_, FalCredentialManager>, + storyteller_creds_manager: State<'_, StorytellerCredentialManager>, +) -> CommandResult { + + info!("fal_background_removal_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 = remove_background( + request, + &app_data_root, + &fal_creds_manager, + &storyteller_creds_manager + ).await; + + match result { + Err(err) => { + 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) { + // _ => {}, + //} + + Err(CommandErrorResponseWrapper { + status, + error_message: Some(error_message.to_string()), + error_type: Some(error_type), + error_details: None, + }) + } + Ok(result) => { + //let event = SoraImageEnqueueSuccessEvent {}; + //let result = event.send(&app); + //if let Err(err) = result { + // error!("Failed to emit event: {:?}", err); + //} + + Ok(FalBackgroundRemovalSuccessResponse { + media_token: result.media_file.token, + cdn_url: result.media_file.media_links.cdn_url, + }.into()) + } + } +} + +#[derive(Debug)] +enum InnerError { + FalError(FalErrorPlus), + AnyhowError(AnyhowError), + StorytellerApiError(ApiError), +} + +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) + } +} + +pub async fn remove_background( + request: FalBackgroundRemovalRequest, + app_data_root: &AppDataRoot, + fal_creds_manager: &FalCredentialManager, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> Result { + + let api_key = fal_creds_manager.get_key_required()?; + let creds = storyteller_creds_manager.get_credentials_required()?; + + let response = get_media_file(&ApiHost::Storyteller, &request.image_media_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); + + simple_http_download(&media_file_url, &filename).await?; + + info!("Calling FAL image background removal..."); + + let result = remove_background_rembg(filename, &api_key).await?; + + let extension_with_dot = get_url_file_extension(&result.image_url) + .map(|ext| format!(".{}", ext)) + .unwrap_or_else(|| ".png".to_string()); + + let result_filename = format!("{}_bg_removed{}", response.media_file.token.as_str(), extension_with_dot); + + simple_http_download(&media_file_url, &result_filename).await?; + + info!("Uploading image file..."); + + let upload_result = upload_image_media_file_from_file( + &ApiHost::Storyteller, + Some(&creds), + result_filename + ).await?; + + let response = get_media_file(&ApiHost::Storyteller, &upload_result.media_file_token).await?; + + Ok(response) +} diff --git a/crates/desktop/artcraft/src/commands/fal/mod.rs b/crates/desktop/artcraft/src/commands/fal/mod.rs index af375b9758..b9da97eefd 100644 --- a/crates/desktop/artcraft/src/commands/fal/mod.rs +++ b/crates/desktop/artcraft/src/commands/fal/mod.rs @@ -1 +1,2 @@ +pub mod fal_background_removal_command; pub mod set_fal_api_key_command; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/commands/sora/sora_image_remix_command.rs b/crates/desktop/artcraft/src/commands/sora/sora_image_remix_command.rs index 54f0608a53..f86eae9e33 100644 --- a/crates/desktop/artcraft/src/commands/sora/sora_image_remix_command.rs +++ b/crates/desktop/artcraft/src/commands/sora/sora_image_remix_command.rs @@ -91,7 +91,7 @@ pub async fn sora_image_remix_command( let has_credentials = sora_creds_manager .has_apparently_complete_credentials() - .unwrap_or(true); + .unwrap_or(true); // NB: Failures would be lock issues if !has_credentials { warn!("No apparently completed credentials found"); diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index 0b3664f650..4e9374f51d 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/crates/desktop/artcraft/src/lib.rs @@ -9,6 +9,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::set_fal_api_key_command::set_fal_api_key_command; use crate::commands::flip_image::flip_image; use crate::commands::platform_info_command::platform_info_command; @@ -143,6 +144,7 @@ pub fn run() { .manage(sora_task_queue) .invoke_handler(tauri::generate_handler![ check_sora_session_command, + fal_background_removal_command, flip_image, get_app_preferences_command, open_sora_login_command, diff --git a/crates/desktop/artcraft/src/state/fal/fal_credential_manager.rs b/crates/desktop/artcraft/src/state/fal/fal_credential_manager.rs index cf43b1ca11..5493645ad4 100644 --- a/crates/desktop/artcraft/src/state/fal/fal_credential_manager.rs +++ b/crates/desktop/artcraft/src/state/fal/fal_credential_manager.rs @@ -49,10 +49,19 @@ impl FalCredentialManager { self.key.get_clone() } + pub fn get_key_required(&self) -> AnyhowResult { + self.key.get_clone_required() + } + pub fn clear_key(&self) -> AnyhowResult<()> { self.key.clear() } + /// Does not check the token for validity + pub fn has_apparent_api_token(&self) -> AnyhowResult { + self.key.is_some() + } + pub fn purge_api_key_from_disk(&self) -> Result<(), io::Error> { let filename = self.app_data_root.credentials_dir().get_fal_api_key_file_path(); if filename.exists() && filename.is_file() { diff --git a/crates/lib/api_clients/fal_client/Cargo.toml b/crates/lib/api_clients/fal_client/Cargo.toml index 34605dae13..91e806b917 100644 --- a/crates/lib/api_clients/fal_client/Cargo.toml +++ b/crates/lib/api_clients/fal_client/Cargo.toml @@ -25,6 +25,7 @@ anyhow = { workspace = true } fal = { version = "0.3.6", features = ["endpoints"] } futures = { workspace = true } tokio = { workspace = true } +url = { workspace = true } [dev-dependencies] testing = { path = "../../testing" } 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 b0de4d4002..5168884a73 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 @@ -8,6 +8,8 @@ pub enum FalErrorPlus { FalError(fal::FalError), /// Another error we didn't handle. AnyhowError(anyhow::Error), + /// URL parse errors. + UrlParseError(url::ParseError), } impl Display for FalErrorPlus { @@ -15,6 +17,7 @@ impl Display for FalErrorPlus { match self { Self::FalError(err) => write!(f, "FalErrorPlus::FalError: {:?}", err), Self::AnyhowError(err) => write!(f, "FalErrorPlus::AnyhowError: {:?}", err), + Self::UrlParseError(err) => write!(f, "FalErrorPlus::UrlParseError: {:?}", err), } } } @@ -32,3 +35,9 @@ impl From for FalErrorPlus { FalErrorPlus::AnyhowError(err) } } + +impl From for FalErrorPlus { + fn from(err: url::ParseError) -> Self { + FalErrorPlus::UrlParseError(err) + } +} diff --git a/crates/lib/api_clients/fal_client/src/requests/remove_background_rembg.rs b/crates/lib/api_clients/fal_client/src/requests/remove_background_rembg.rs index 27dd7d8ce4..d04c77a838 100644 --- a/crates/lib/api_clients/fal_client/src/requests/remove_background_rembg.rs +++ b/crates/lib/api_clients/fal_client/src/requests/remove_background_rembg.rs @@ -6,6 +6,7 @@ use fal::prelude::*; use fal_client::file_to_base64_url::file_to_base64_url; use futures::StreamExt; use std::path::Path; +use url::Url; pub struct RemBgArgs<'a, P: AsRef> { pub image_path: P, @@ -13,12 +14,12 @@ pub struct RemBgArgs<'a, P: AsRef> { } pub struct RemBgResponse { - pub image_url: String, + pub image_url: Url, } /// remove background using `fal_ai/imageutils/rembg` -pub async fn remove_background_rembg>(args: RemBgArgs<'_, P>) -> Result { - let image_url = file_to_base64_url(&args.image_path)?; +pub async fn remove_background_rembg>(image_path: P, api_key: &FalApiKey) -> Result { + let image_url = file_to_base64_url(image_path)?; let request = RemoveBackgroundInput { image_url, @@ -27,7 +28,7 @@ pub async fn remove_background_rembg>(args: RemBgArgs<'_, P>) -> }; let result = rembg(request) - .with_api_key(&args.api_key.0) + .with_api_key(&api_key.0) .queue() .await?; @@ -42,15 +43,17 @@ pub async fn remove_background_rembg>(args: RemBgArgs<'_, P>) -> let output = result.response().await?; + let url = Url::parse(&output.image.url)?; + Ok(RemBgResponse { - image_url: output.image.url, + image_url: url, }) } #[cfg(test)] mod tests { use crate::creds::fal_api_key::FalApiKey; - use crate::requests::remove_background_rembg::{remove_background_rembg, RemBgArgs}; + use crate::requests::remove_background_rembg::remove_background_rembg; use errors::AnyhowResult; use testing::test_file_path::test_file_path; @@ -63,10 +66,7 @@ mod tests { const KEY : &str = "DO NOT COMMIT"; let api_key = FalApiKey::from_str(KEY); - let result = remove_background_rembg(RemBgArgs { - image_path: image, - api_key: &api_key, - }).await?; + let result = remove_background_rembg(image, &api_key).await?; Ok(()) } diff --git a/crates/lib/caching/memory_store/src/clone_slot.rs b/crates/lib/caching/memory_store/src/clone_slot.rs index 3dd13c3742..ef2a3d1001 100644 --- a/crates/lib/caching/memory_store/src/clone_slot.rs +++ b/crates/lib/caching/memory_store/src/clone_slot.rs @@ -31,7 +31,17 @@ impl CloneSlot { Ok(store) => Ok(store.clone()), } } - + + pub fn get_clone_required(&self) -> AnyhowResult { + match self.store.read() { + Err(err) => Err(anyhow!("Failed to acquire read lock: {:?}", err)), + Ok(store) => match &*store { + None => Err(anyhow!("Required stored item is not set")), + Some(inner) => Ok(inner.clone()), + } + } + } + pub fn set_clone(&self, value: &T) -> AnyhowResult<()> { match self.store.write() { Err(err) => Err(anyhow!("Failed to acquire write lock: {:?}", err)),