From f5af974b6c0dd960abc844ccbc1e64a21e4068cb Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 14 Jul 2025 21:05:42 -0400 Subject: [PATCH] support aspect ratios in text-to-image --- .../image/enqueue_text_to_image_command.rs | 27 ++++++++-- .../image/generic/handle_image_artcraft.rs | 50 +++++++++++++++---- .../image/gpt_image_1/handle_image_sora.rs | 12 ++++- 3 files changed, 73 insertions(+), 16 deletions(-) diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image/enqueue_text_to_image_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image/enqueue_text_to_image_command.rs index 5d54e00255..8398efcc75 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image/enqueue_text_to_image_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image/enqueue_text_to_image_command.rs @@ -36,8 +36,25 @@ pub struct EnqueueTextToImageRequest { /// The model to use. pub model: Option, + + /// Aspect ratio. + pub aspect_ratio: Option, } +// TODO(bt,2025-07-14): Support other aspect ratios / resolutions - +// Flux Dev has: 4:3, 16:9, 3:4, 9:16, and custom. +// Flux Schnell has: 4:3, 16:9, 3:4, 9:16, and custom. +// Flux Pro has: 4:3, 16:9, 3:4, 9:16, and custom. +// Flux Pro Ulra has: 4:3, 16:9, 3:4, 9:16, and custom. + +#[derive(Deserialize, Debug, Copy, Clone)] +#[serde(rename_all = "snake_case")] +pub enum TextToImageSize { + Auto, + Square, + Wide, + Tall, +} #[derive(Serialize)] pub struct EnqueueTextToImageSuccessResponse { @@ -65,8 +82,8 @@ pub enum EnqueueTextToImageErrorType { #[tauri::command] pub async fn enqueue_text_to_image_command( - app: AppHandle, request: EnqueueTextToImageRequest, + app: AppHandle, app_data_root: State<'_, AppDataRoot>, app_env_configs: State<'_, AppEnvConfigs>, provider_priority_store: State<'_, ProviderPriorityStore>, @@ -81,8 +98,8 @@ pub async fn enqueue_text_to_image_command( info!("enqueue_text_to_image called"); let result = handle_request( - &app, request, + &app, &app_data_root, &provider_priority_store, &task_database, @@ -146,8 +163,8 @@ pub async fn enqueue_text_to_image_command( pub async fn handle_request( - app: &AppHandle, request: EnqueueTextToImageRequest, + app: &AppHandle, app_data_root: &AppDataRoot, provider_priority_store: &ProviderPriorityStore, task_database: &TaskDatabase, @@ -160,8 +177,8 @@ pub async fn handle_request( ) -> Result { let result = dispatch_request( - &app, request, + &app, &app_data_root, &provider_priority_store, &fal_creds_manager, @@ -190,8 +207,8 @@ pub async fn handle_request( } pub async fn dispatch_request( - app: &AppHandle, request: EnqueueTextToImageRequest, + app: &AppHandle, app_data_root: &AppDataRoot, provider_priority_store: &ProviderPriorityStore, fal_creds_manager: &FalCredentialManager, diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image/generic/handle_image_artcraft.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image/generic/handle_image_artcraft.rs index 4db69aa4f8..d9f32acecc 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image/generic/handle_image_artcraft.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image/generic/handle_image_artcraft.rs @@ -1,4 +1,4 @@ -use crate::core::commands::enqueue::image::enqueue_text_to_image_command::EnqueueTextToImageRequest; +use crate::core::commands::enqueue::image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize}; use crate::core::commands::enqueue::image::internal_image_error::InternalImageError; use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; @@ -9,10 +9,10 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; use anyhow::anyhow; -use artcraft_api_defs::generate::image::generate_flux_1_dev_text_to_image::GenerateFlux1DevTextToImageRequest; -use artcraft_api_defs::generate::image::generate_flux_1_schnell_text_to_image::GenerateFlux1SchnellTextToImageRequest; -use artcraft_api_defs::generate::image::generate_flux_pro_11_text_to_image::GenerateFluxPro11TextToImageRequest; -use artcraft_api_defs::generate::image::generate_flux_pro_11_ultra_text_to_image::GenerateFluxPro11UltraTextToImageRequest; +use artcraft_api_defs::generate::image::generate_flux_1_dev_text_to_image::{GenerateFlux1DevTextToImageAspectRatio, GenerateFlux1DevTextToImageRequest}; +use artcraft_api_defs::generate::image::generate_flux_1_schnell_text_to_image::{GenerateFlux1SchnellTextToImageAspectRatio, GenerateFlux1SchnellTextToImageRequest}; +use artcraft_api_defs::generate::image::generate_flux_pro_11_text_to_image::{GenerateFluxPro11TextToImageAspectRatio, GenerateFluxPro11TextToImageRequest}; +use artcraft_api_defs::generate::image::generate_flux_pro_11_ultra_text_to_image::{GenerateFluxPro11UltraTextToImageAspectRatio, GenerateFluxPro11UltraTextToImageRequest}; use enums::common::generation_provider::GenerationProvider; use enums::tauri::tasks::task_type::TaskType; use fal_client::requests::queue::image_gen::enqueue_flux_pro_11_ultra_text_to_image::{enqueue_flux_pro_11_ultra_text_to_image, FluxPro11UltraTextToImageArgs}; @@ -66,10 +66,18 @@ pub async fn handle_image_artcraft( Some(ImageModel::Flux1Dev) => { info!("enqueue Flux 1 Dev"); selected_model = GenerationModel::Flux1Dev; + let aspect_ratio = request.aspect_ratio + .map(|aspect| match aspect { + // TODO(bt,2025-07-14): Support other aspect ratios. + TextToImageSize::Tall => GenerateFlux1DevTextToImageAspectRatio::PortraitNineBySixteen, + TextToImageSize::Wide => GenerateFlux1DevTextToImageAspectRatio::LandscapeSixteenByNine, + TextToImageSize::Auto => GenerateFlux1DevTextToImageAspectRatio::SquareHd, + TextToImageSize::Square => GenerateFlux1DevTextToImageAspectRatio::SquareHd, + }); let request = GenerateFlux1DevTextToImageRequest { uuid_idempotency_token, prompt: request.prompt, - aspect_ratio: None, + aspect_ratio, num_images: None, }; let result = generate_flux_1_dev_text_to_image( @@ -91,10 +99,18 @@ pub async fn handle_image_artcraft( Some(ImageModel::Flux1Schnell) => { info!("enqueue Flux 1 Schnell"); selected_model = GenerationModel::Flux1Schnell; + let aspect_ratio = request.aspect_ratio + .map(|aspect| match aspect { + // TODO(bt,2025-07-14): Support other aspect ratios. + TextToImageSize::Tall => GenerateFlux1SchnellTextToImageAspectRatio::PortraitNineBySixteen, + TextToImageSize::Wide => GenerateFlux1SchnellTextToImageAspectRatio::LandscapeSixteenByNine, + TextToImageSize::Auto => GenerateFlux1SchnellTextToImageAspectRatio::SquareHd, + TextToImageSize::Square => GenerateFlux1SchnellTextToImageAspectRatio::SquareHd, + }); let request = GenerateFlux1SchnellTextToImageRequest { uuid_idempotency_token, prompt: request.prompt, - aspect_ratio: None, + aspect_ratio, num_images: None, }; let result = generate_flux_1_schnell_text_to_image( @@ -116,10 +132,18 @@ pub async fn handle_image_artcraft( Some(ImageModel::FluxPro11) => { info!("enqueue Flux Pro 1.1"); selected_model = GenerationModel::FluxPro11; + let aspect_ratio = request.aspect_ratio + .map(|aspect| match aspect { + // TODO(bt,2025-07-14): Support other aspect ratios. + TextToImageSize::Tall => GenerateFluxPro11TextToImageAspectRatio::PortraitNineBySixteen, + TextToImageSize::Wide => GenerateFluxPro11TextToImageAspectRatio::LandscapeSixteenByNine, + TextToImageSize::Auto => GenerateFluxPro11TextToImageAspectRatio::Square, + TextToImageSize::Square => GenerateFluxPro11TextToImageAspectRatio::Square, + }); let request = GenerateFluxPro11TextToImageRequest { uuid_idempotency_token, prompt: request.prompt, - aspect_ratio: None, + aspect_ratio, num_images: None, }; let result = generate_flux_pro_11_text_to_image( @@ -141,10 +165,18 @@ pub async fn handle_image_artcraft( Some(ImageModel::FluxPro11Ultra) => { info!("enqueue Flux Pro 1.1 Ultra"); selected_model = GenerationModel::FluxPro11Ultra; + let aspect_ratio = request.aspect_ratio + .map(|aspect| match aspect { + // TODO(bt,2025-07-14): Support other aspect ratios. + TextToImageSize::Tall => GenerateFluxPro11UltraTextToImageAspectRatio::PortraitNineBySixteen, + TextToImageSize::Wide => GenerateFluxPro11UltraTextToImageAspectRatio::LandscapeSixteenByNine, + TextToImageSize::Auto => GenerateFluxPro11UltraTextToImageAspectRatio::Square, + TextToImageSize::Square => GenerateFluxPro11UltraTextToImageAspectRatio::Square, + }); let request = GenerateFluxPro11UltraTextToImageRequest { uuid_idempotency_token, prompt: request.prompt, - aspect_ratio: None, + aspect_ratio, num_images: None, }; let result = generate_flux_pro_11_ultra_text_to_image( diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image/gpt_image_1/handle_image_sora.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image/gpt_image_1/handle_image_sora.rs index a8f6d1484e..34a206843b 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image/gpt_image_1/handle_image_sora.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image/gpt_image_1/handle_image_sora.rs @@ -1,4 +1,4 @@ -use crate::core::commands::enqueue::image::enqueue_text_to_image_command::EnqueueTextToImageRequest; +use crate::core::commands::enqueue::image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize}; use crate::core::commands::enqueue::image::internal_image_error::InternalImageError; use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; @@ -46,12 +46,20 @@ pub async fn handle_image_sora( } let prompt = request.prompt.as_deref().unwrap_or(""); + + let size = match request.aspect_ratio { + None => ImageSize::Square, + Some(TextToImageSize::Auto) => ImageSize::Square, + Some(TextToImageSize::Square) => ImageSize::Square, + Some(TextToImageSize::Tall) => ImageSize::Tall, + Some(TextToImageSize::Wide) => ImageSize::Wide, + }; let result = simple_image_gen_with_session_auto_renew(SimpleImageGenAutoRenewRequest { prompt: prompt.to_string(), num_images: NumImages::One, - image_size: ImageSize::Square, + image_size: size, credentials: &creds, request_timeout: Some(SORA_SIMPLE_IMAGE_GEN_TIMEOUT), }).await;