From d60e123aba910d7a5c58207e5df9d3a73471a65f Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 26 May 2025 19:37:48 -0400 Subject: [PATCH] enqueue tasks --- .../enqueue/enqueue_text_to_image_command.rs | 30 ++++++++++++++----- .../app/src/pages/PageImage/TextToImage.tsx | 10 +++---- .../promptbox/src/lib/PromptBoxImage.tsx | 26 ++++++++++++---- frontend/libs/tauri-api/src/index.ts | 1 + 4 files changed, 48 insertions(+), 19 deletions(-) diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/enqueue_text_to_image_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/enqueue_text_to_image_command.rs index a9fddce81d..78ecc54f39 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/enqueue_text_to_image_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/enqueue_text_to_image_command.rs @@ -53,6 +53,7 @@ use storyteller_client::utils::api_host::ApiHost; use tauri::{AppHandle, Emitter, Manager, State}; use tempfile::NamedTempFile; use tokens::tokens::media_files::MediaFileToken; +use crate::services::fal::state::fal_task_queue::FalTaskQueue; #[derive(Deserialize)] pub struct EnqueueTextToImageRequest { @@ -98,6 +99,7 @@ pub async fn enqueue_text_to_image_command( app_data_root: State<'_, AppDataRoot>, fal_creds_manager: State<'_, FalCredentialManager>, storyteller_creds_manager: State<'_, StorytellerCredentialManager>, + fal_task_queue: State<'_, FalTaskQueue>, ) -> CommandResult { info!("enqueue_text_to_image called"); @@ -106,7 +108,8 @@ pub async fn enqueue_text_to_image_command( request, &app_data_root, &fal_creds_manager, - &storyteller_creds_manager + &storyteller_creds_manager, + &fal_task_queue, ).await; match result { @@ -203,6 +206,7 @@ pub async fn handle_request( app_data_root: &AppDataRoot, fal_creds_manager: &FalCredentialManager, storyteller_creds_manager: &StorytellerCredentialManager, + fal_task_queue: &FalTaskQueue, ) -> Result<(), InnerError> { match request.model { @@ -222,8 +226,6 @@ pub async fn handle_request( } let api_key = fal_creds_manager.get_key_required()?; - let creds = storyteller_creds_manager.get_credentials_required()?; - let prompt = request.prompt.as_deref().unwrap_or(""); let result = match request.model { @@ -231,24 +233,36 @@ pub async fn handle_request( return Err(InnerError::NoModelSpecified); } Some(EnqueueTextToImageModel::FluxProUltra) => { + info!("enqueue Flux Pro Ultra text-to-image with prompt: {}", prompt); enqueue_flux_pro_ultra_text_to_image(FluxProUltraTextToImageArgs { prompt, api_key: &api_key, }).await } Some(EnqueueTextToImageModel::Recraft3) => { + info!("enqueue Recraft v3 text-to-image with prompt: {}", prompt); enqueue_recraft3_text_to_image(Recraft3TextToImageArgs { prompt, api_key: &api_key, }).await } }; - - if let Err(err) = result { - error!("Error enqueuing text to image: {:?}", err); - return Err(InnerError::FalError(err)); + + match result { + Ok(enqueued) => { + info!("Successfully enqueued text to image"); + + if let Err(err) = fal_task_queue.insert(&enqueued) { + error!("Failed to enqueue task: {:?}", err); + return Err(InnerError::AnyhowError(anyhow!("Failed to enqueue task: {:?}", err))); + } + } + Err(err) => { + error!("Failed to enqueue text to image: {:?}", err); + return Err(InnerError::FalError(err)); + } } - + Ok(()) } diff --git a/frontend/apps/editor3d/app/src/pages/PageImage/TextToImage.tsx b/frontend/apps/editor3d/app/src/pages/PageImage/TextToImage.tsx index f72d0cd238..56a75b9fca 100644 --- a/frontend/apps/editor3d/app/src/pages/PageImage/TextToImage.tsx +++ b/frontend/apps/editor3d/app/src/pages/PageImage/TextToImage.tsx @@ -13,7 +13,7 @@ interface TextToImageProps { const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => { const containerRef = useRef(null); - const [selectedModel, setSelectedModel] = useState("Kling 2.0"); + const [selectedModel, setSelectedModel] = useState("Flux Pro Ultra"); const [useSystemPrompt, setUseSystemPrompt] = useState(false); const [isImageBeingProcessed, setIsImageBeingProcessed] = useState(false); const [selectedImage, setSelectedImage] = useState(null); @@ -29,16 +29,16 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => { const modelList: PopoverItem[] = [ { - label: "Kling 2.0", + label: "Flux Pro Ultra", icon: , - selected: selectedModel === "Kling 2.0", + selected: selectedModel === "Flux Pro Ultra", description: "High quality model", badges: [{ label: "2 min.", icon: }], }, { - label: "Runway", + label: "Recraft 3", icon: , - selected: selectedModel === "Runway", + selected: selectedModel === "Recraft 3", description: "Fast and high-quality model", badges: [{ label: "2 min.", icon: }], }, diff --git a/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx b/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx index afe1d68f6e..25f7233c22 100644 --- a/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx +++ b/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx @@ -7,7 +7,7 @@ import { PopoverMenu, PopoverItem } from "@storyteller/ui-popover"; import { Tooltip } from "@storyteller/ui-tooltip"; import { Button, ToggleButton } from "@storyteller/ui-button"; import { Modal } from "@storyteller/ui-modal"; -import { FalKlingImageToVideo } from "@storyteller/tauri-api"; +import { EnqueueTextToImage, EnqueueTextToImageModel } from "@storyteller/tauri-api"; import { faMessageXmark, faMessageCheck, @@ -71,8 +71,8 @@ export const PromptBoxImage = ({ console.log("Is this a desktop app?", IsDesktopApp()); const [isModalOpen, setIsModalOpen] = useState(false); const [content, setContent] = useState(""); - const { jobTokens, addJobToken, removeJobToken, clearJobTokens } = - useJobContext(); + //const { jobTokens, addJobToken, removeJobToken, clearJobTokens } = + // useJobContext(); const [prompt, setPrompt] = useState(""); const [isEnqueueing, setIsEnqueueing] = useState(false); @@ -199,12 +199,15 @@ export const PromptBoxImage = ({ const handleEnqueue = async () => { setIsEnqueueing(true); - console.log("generateResponse", generateResponse); + const modelName = getModelByName(model); - const generateResponse = await FalKlingImageToVideo({ - image_media_token: referenceImages[0].mediaToken, + const generateResponse = await EnqueueTextToImage({ + prompt: prompt, + model: modelName, }); + console.log("generateResponse", generateResponse); + onEnqueuePressed?.(); if (generateResponse.status === CommandSuccessStatus.Success) { const prefs = await GetAppPreferences(); @@ -259,6 +262,17 @@ export const PromptBoxImage = ({ return iconElement.props.icon; }; + const getModelByName = (name: string) : EnqueueTextToImageModel => { + switch (name) { + case "Flux Pro Ultra": + return EnqueueTextToImageModel.FluxProUltra; + case "Recraft 3": + return EnqueueTextToImageModel.Recraft3; + default: + return EnqueueTextToImageModel.FluxProUltra; + } + }; + const handleKeyDown = (e: React.KeyboardEvent) => { if (e.key === "Enter" && !e.shiftKey) { e.preventDefault(); diff --git a/frontend/libs/tauri-api/src/index.ts b/frontend/libs/tauri-api/src/index.ts index f4112aeee0..f3c5b1ea38 100644 --- a/frontend/libs/tauri-api/src/index.ts +++ b/frontend/libs/tauri-api/src/index.ts @@ -1,4 +1,5 @@ export * from "./lib/common/CommandStatus.js"; +export * from "./lib/enqueue/EnqueueTextToImage.js"; export * from "./lib/fal/FalBackgroundRemoval.js"; export * from "./lib/fal/FalHunyuanImageTo3d.js"; export * from "./lib/fal/FalKlingImageToVideo.js";