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 6612af499f..5882fb2b34 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 @@ -12,14 +12,17 @@ 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, ImageReader}; +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}; @@ -33,11 +36,10 @@ 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, Write}; +use std::io::{Cursor, Read, Write}; use std::ops::Add; use std::path::PathBuf; use std::time::Duration; -use anyhow::anyhow; 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}; @@ -45,9 +47,7 @@ use storyteller_client::media_files::upload_image_media_file_from_file::upload_i use storyteller_client::utils::api_host::ApiHost; use tauri::{AppHandle, Emitter, Manager, State}; use tempfile::NamedTempFile; -use mimetypes::mimetype_info::mimetype_info::MimetypeInfo; use tokens::tokens::media_files::MediaFileToken; -use crate::utils::simple_http_download_to_tempfile::simple_http_download_to_tempfile; #[derive(Deserialize)] pub struct FalBackgroundRemovalRequest { @@ -64,6 +64,8 @@ pub struct FalBackgroundRemovalSuccessResponse { pub media_token: MediaFileToken, /// Result URL pub cdn_url: Url, + /// Base64 bytes + pub base64_bytes: String, } impl SerializeMarker for FalBackgroundRemovalSuccessResponse {} @@ -92,6 +94,11 @@ pub async fn fal_background_removal_command( .has_apparent_api_token() .unwrap_or(true); // NB: Failures would be lock issues + let maybe_key = fal_creds_manager.get_key().expect("wat") + .map(|key| key.0); + + println!(">>> fal_background_removal_command; has_creds: {} maybe_key: {:?}", has_credentials, maybe_key); + if !has_credentials { warn!("No API key found"); return Err(CommandErrorResponseWrapper { @@ -142,8 +149,9 @@ pub async fn fal_background_removal_command( //} Ok(FalBackgroundRemovalSuccessResponse { - media_token: result.media_file.token, - cdn_url: result.media_file.media_links.cdn_url, + media_token: result.0.media_file.token, + cdn_url: result.0.media_file.media_links.cdn_url, + base64_bytes: result.1, }.into()) } } @@ -200,7 +208,7 @@ pub async fn remove_background( app_data_root: &AppDataRoot, fal_creds_manager: &FalCredentialManager, storyteller_creds_manager: &StorytellerCredentialManager, -) -> Result { +) -> Result<(GetMediaFileSuccessResponse, String), InnerError> { let api_key = fal_creds_manager.get_key_required()?; let creds = storyteller_creds_manager.get_credentials_required()?; @@ -245,7 +253,12 @@ pub async fn remove_background( let response = get_media_file(&ApiHost::Storyteller, &upload_result.media_file_token).await?; - Ok(response) + let mut buffer = Vec::new(); + temp_download.read_to_end(&mut buffer)?; + + let bytes = BASE64_STANDARD.encode(buffer.as_bytes()); + + Ok((response, bytes)) } async fn download_media_file(app_data_root: &AppDataRoot, token: &MediaFileToken) -> Result { diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index 4e9374f51d..1c4171dd3a 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/crates/desktop/artcraft/src/lib.rs @@ -50,6 +50,7 @@ pub fn run() { println!("Attempting to read existing artcraft credentials..."); let storyteller_creds_manager = StorytellerCredentialManager::initialize_from_disk_infallible(&app_data_root); let storyteller_creds_manager_2 = storyteller_creds_manager.clone(); + let storyteller_creds_manager_3 = storyteller_creds_manager.clone(); println!("Attempting to read existing sora credentials..."); let sora_creds_manager = SoraCredentialManager::initialize_from_disk_infallible(&app_data_root); @@ -142,6 +143,7 @@ pub fn run() { .manage(sora_creds_manager) .manage(fal_creds_manager) .manage(sora_task_queue) + .manage(storyteller_creds_manager_3) .invoke_handler(tauri::generate_handler![ check_sora_session_command, fal_background_removal_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 5493645ad4..95ba34c74b 100644 --- a/crates/desktop/artcraft/src/state/fal/fal_credential_manager.rs +++ b/crates/desktop/artcraft/src/state/fal/fal_credential_manager.rs @@ -36,7 +36,7 @@ impl FalCredentialManager { } Self { - key: CloneSlot::empty(), + key: holder, app_data_root: app_data_root.clone(), } } diff --git a/frontend/apps/editor3d/app/src/pages/Page2d/Engine.ts b/frontend/apps/editor3d/app/src/pages/Page2d/Engine.ts index 771bc8c81f..3b6c90ed99 100644 --- a/frontend/apps/editor3d/app/src/pages/Page2d/Engine.ts +++ b/frontend/apps/editor3d/app/src/pages/Page2d/Engine.ts @@ -28,6 +28,7 @@ import { invoke } from "@tauri-apps/api/core"; import { DecodeBase64ToImage } from "./utilities/DecodeBase64ToImage"; import { EncodeImageBitmapToBase64 } from "./utilities/EncodeImageBitmapToBase64"; +import { FalBackgroundRemoval, FalBackgroundRemovalRequest } from "@storyteller/tauri-api"; @@ -443,6 +444,8 @@ export class Engine { } removeBackground() { + console.log("Remove background"); + const selectedNodeSet = this.selectionManager.getSelectedNodes(); if (selectedNodeSet.size !== 1) { @@ -460,15 +463,30 @@ export class Engine { ) => { const image: HTMLImageElement = selectedNode.kNode.image(); const base64Bitmap = await EncodeImageBitmapToBase64(image); - const base64BitmapResponse = await invoke("remove_background", { - image: base64Bitmap, - }); - const createdImage = await DecodeBase64ToImage( - base64BitmapResponse as string, - ); + const falRemoveBgRequest: FalBackgroundRemovalRequest = { + base64_image: base64Bitmap, + }; - callback(selectedNode, createdImage); + const response = await FalBackgroundRemoval(falRemoveBgRequest); + + if (response.status !== "success") { + console.error("Failed to remove background", response); + // TODO + return; + } + + console.log("Remove background success", response) + + if ("payload" in response) { + const base64BitmapResponse = response.payload?.base64_bytes; + + const createdImage = await DecodeBase64ToImage( + base64BitmapResponse as string, + ); + + callback(selectedNode, createdImage); + } }; const _promise = removeBgInvocation( diff --git a/frontend/libs/tauri-api/src/index.ts b/frontend/libs/tauri-api/src/index.ts index 56c11deb4a..dd2e9811c2 100644 --- a/frontend/libs/tauri-api/src/index.ts +++ b/frontend/libs/tauri-api/src/index.ts @@ -5,3 +5,4 @@ export * from "./lib/sora/CheckSoraSession.js"; export * from "./lib/sora/LogoutSoraSession.js"; export * from "./lib/sora/SoraImageRemix.js"; export * from "./lib/sora/useSoraLoginListener.js"; +export * from "./lib/fal/FalBackgroundRemoval.js"; diff --git a/frontend/libs/tauri-api/src/lib/fal/FalBackgroundRemoval.ts b/frontend/libs/tauri-api/src/lib/fal/FalBackgroundRemoval.ts new file mode 100644 index 0000000000..f92ca53269 --- /dev/null +++ b/frontend/libs/tauri-api/src/lib/fal/FalBackgroundRemoval.ts @@ -0,0 +1,39 @@ +import { invoke } from "@tauri-apps/api/core"; +import { CommandResult } from "../common/CommandStatus"; + +export enum FalBackgroundRemovalErrorType { + ServerError = "server_error", + NeedsFalApiKey = "needs_fal_api_key", +} + +export interface FalBackgroundRemovalRequest { + image_media_token?: string; + base64_image?: string; +} + +export interface FalBackgroundRemovalError extends CommandResult { + error_type: FalBackgroundRemovalErrorType; + error_message?: string; +} + +export interface FalBackgroundRemovalPayload { + media_token: string; + cdn_url: string; + base64_bytes: string; +} + +export interface FalBackgroundRemovalSuccess extends CommandResult { + payload: FalBackgroundRemovalPayload; +} + +export type FalBackgroundRemovalResult = FalBackgroundRemovalSuccess | FalBackgroundRemovalError; + +export const FalBackgroundRemoval = async (request: FalBackgroundRemovalRequest) : Promise => { + let result = await invoke("fal_background_removal_command", { + request: { + image_media_token: request.image_media_token, + base64_image: request.base64_image, + } + }); + return (result as FalBackgroundRemovalResult); +}