fal background removal endpoint

Cargo.lock

@@ -4254,6 +4254,7 @@ dependencies = [
  "futures 0.3.31",
  "testing",
  "tokio",
+ "url",
  "workspace-build-acceleration",
 ]
 

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<()> {

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<FalBackgroundRemovalSuccessResponse, SoraImageRemixErrorType, ()> {
+  
+  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<AnyhowError> for InnerError {
+  fn from(value: AnyhowError) -> Self {
+    Self::AnyhowError(value)
+  }
+}
+
+impl From<FalErrorPlus> for InnerError {
+  fn from(value: FalErrorPlus) -> Self {
+    Self::FalError(value)
+  }
+}
+
+impl From<ApiError> 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<GetMediaFileSuccessResponse, InnerError> {
+  
+  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)
+}

crates/desktop/artcraft/src/commands/fal/mod.rs

@@ -1 +1,2 @@
+pub mod fal_background_removal_command;
 pub mod set_fal_api_key_command;

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");

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,

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<FalApiKey> {
+    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<bool> {
+    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() {

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" }

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<anyhow::Error> for FalErrorPlus {
     FalErrorPlus::AnyhowError(err)
   }
 }
+
+impl From<url::ParseError> for FalErrorPlus {
+  fn from(err: url::ParseError) -> Self {
+    FalErrorPlus::UrlParseError(err)
+  }
+}

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<Path>> {
   pub image_path: P,
@@ -13,12 +14,12 @@ pub struct RemBgArgs<'a, P: AsRef<Path>> {
 }
 
 pub struct RemBgResponse {
-  pub image_url: String,
+  pub image_url: Url,
 }
 
 /// remove background using `fal_ai/imageutils/rembg`
-pub async fn remove_background_rembg<P: AsRef<Path>>(args: RemBgArgs<'_, P>) -> Result<RemBgResponse, FalErrorPlus> {
-  let image_url = file_to_base64_url(&args.image_path)?;
+pub async fn remove_background_rembg<P: AsRef<Path>>(image_path: P, api_key: &FalApiKey) -> Result<RemBgResponse, FalErrorPlus> {
+  let image_url = file_to_base64_url(image_path)?;
   
   let request = RemoveBackgroundInput {
     image_url,
@@ -27,7 +28,7 @@ pub async fn remove_background_rembg<P: AsRef<Path>>(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<P: AsRef<Path>>(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(())
   }

crates/lib/caching/memory_store/src/clone_slot.rs

@@ -31,7 +31,17 @@ impl <T: Clone> CloneSlot<T> {
       Ok(store) => Ok(store.clone()),
     }
   }
-  
+
+  pub fn get_clone_required(&self) -> AnyhowResult<T> {
+    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)),