From 5f18e9f2dfcf220f42e2f307ebffd067172640da Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 27 Nov 2025 02:20:07 -0500 Subject: [PATCH] plumb aspect ratio --- .../messages/websocket_client_message.rs | 33 +------------------ .../enqueue/text_to_image/grok/handle_grok.rs | 12 ++++++- .../grok/state/grok_image_prompt_queue.rs | 4 +++ .../grok_image_websocket_thread.rs | 4 +-- 4 files changed, 18 insertions(+), 35 deletions(-) diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/messages/websocket_client_message.rs b/crates/api_clients/grok_client/src/requests/image_websocket/messages/websocket_client_message.rs index 88941f69e9..f6422fad2c 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/messages/websocket_client_message.rs +++ b/crates/api_clients/grok_client/src/requests/image_websocket/messages/websocket_client_message.rs @@ -49,7 +49,7 @@ pub struct ClientMessageItemContentProperties { pub aspect_ratio: ClientMessageAspectRatio, } -#[derive(Serialize, Clone, Copy)] +#[derive(Serialize, Clone, Copy, Debug)] pub enum ClientMessageAspectRatio { #[serde(rename = "2:3")] TallTwoByThree, @@ -64,37 +64,6 @@ pub enum ClientMessageAspectRatio { impl WebsocketClientMessage { /// Create a new image prompt websocket client message pub fn new_image_prompt(prompt: &str, aspect_ratio: ClientMessageAspectRatio) -> Self { - /* - { - "type":"conversation.item.create", - "timestamp":1764222761003, - "item": { - "type":"message", - "content": [ - { - "requestId":"efdde1fc-2427-4f49-bd90-3e54df0e6294", - "text":"fighter jet in the taco bell drive thru", - "type":"input_text", - "properties": { - "section_count":0, - "is_kids_mode":false, - "enable_nsfw":true, - "skip_upsampler":false, - "is_initial":false, - "aspect_ratio":"2:3" - } - } - ] - } - } - - "2:3" - "3:2" - "1:1" - - {"type":"conversation.item.create","timestamp":1764223129572,"item":{"type":"message","content":[{"requestId":"3825ffdd-0613-48ba-9a50-fb75bbfbdba0","text":"tank at mcdonalds","type":"input_text","properties":{"section_count":0,"is_kids_mode":false,"enable_nsfw":true,"skip_upsampler":false,"is_initial":false,"aspect_ratio":"1:1"}}]}} - {"type":"conversation.item.create","timestamp":1764223055299,"item":{"type":"message","content":[{"requestId":"86361ecf-bdcf-4159-bfcb-f1097dc7499d","text":"an attack helicopter at Burger King","type":"input_scroll","properties":{"section_count":0,"is_kids_mode":false,"enable_nsfw":true,"skip_upsampler":false,"is_initial":false,"aspect_ratio":"3:2"}}]}} - */ Self { r#type: "conversation.item.create".to_string(), timestamp: Utc::now().timestamp_millis() as u64, diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs index c00c03eac6..249f78aeec 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs @@ -1,12 +1,13 @@ use crate::core::commands::enqueue::generate_error::GenerateError; use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; -use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::EnqueueTextToImageRequest; +use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize}; use crate::core::events::generation_events::common::GenerationModel; use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; use crate::services::grok::state::grok_image_prompt_queue::{GrokImagePromptQueue, PromptItem}; use enums::common::generation_provider::GenerationProvider; use enums::tauri::tasks::task_type::TaskType; +use grok_client::requests::image_websocket::messages::websocket_client_message::ClientMessageAspectRatio; use idempotency::uuid::generate_random_uuid; use tauri::AppHandle; @@ -53,9 +54,18 @@ pub async fn handle_grok( .map(|prompt| prompt.trim().to_string()) .unwrap_or_else(|| "".to_string()); + let aspect_ratio = match request.aspect_ratio { + None => ClientMessageAspectRatio::Square, + Some(TextToImageSize::Auto) => ClientMessageAspectRatio::Square, + Some(TextToImageSize::Square) => ClientMessageAspectRatio::Square, + Some(TextToImageSize::Wide) => ClientMessageAspectRatio::WideThreeByTwo, + Some(TextToImageSize::Tall) => ClientMessageAspectRatio::TallTwoByThree, + }; + grok_image_prompt_queue.enqueue(PromptItem { task_id: job_id.clone(), prompt, + aspect_ratio, })?; Ok(TaskEnqueueSuccess { diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_image_prompt_queue.rs b/crates/desktop/artcraft/src/services/grok/state/grok_image_prompt_queue.rs index 6bf31824f4..6e185aa2d9 100644 --- a/crates/desktop/artcraft/src/services/grok/state/grok_image_prompt_queue.rs +++ b/crates/desktop/artcraft/src/services/grok/state/grok_image_prompt_queue.rs @@ -2,6 +2,7 @@ use crate::core::artcraft_error::ArtcraftError; use log::error; use std::collections::VecDeque; use std::sync::{Arc, Mutex}; +use grok_client::requests::image_websocket::messages::websocket_client_message::ClientMessageAspectRatio; #[derive(Clone)] pub struct GrokImagePromptQueue { @@ -16,6 +17,9 @@ pub struct PromptItem { /// Text prompt pub prompt: String, + + /// The aspect ratio of the image + pub aspect_ratio: ClientMessageAspectRatio, } impl GrokImagePromptQueue { diff --git a/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs index 5401461265..4c065d4b2b 100644 --- a/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs +++ b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs @@ -22,6 +22,7 @@ use grok_client::requests::image_websocket::create_listen_websocket::{create_lis use grok_client::requests::image_websocket::grok_websocket::GrokWebsocket; use grok_client::requests::image_websocket::grok_wrapped_websocket::GrokWrappedWebsocket; use grok_client::requests::image_websocket::listen_for_websocket_images::{listen_for_websocket_images, ImageResults, ListenForWebsocketImagesArgs}; +use grok_client::requests::image_websocket::messages::websocket_client_message::ClientMessageAspectRatio; use grok_client::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; use idempotency::uuid::generate_random_uuid; use log::{error, info, warn}; @@ -39,7 +40,6 @@ use storyteller_client::endpoints::prompts::create_prompt::create_prompt; use storyteller_client::error::api_error::ApiError; use storyteller_client::error::storyteller_error::StorytellerError; use tauri::AppHandle; -use grok_client::requests::image_websocket::messages::websocket_client_message::ClientMessageAspectRatio; use tokens::tokens::batch_generations::BatchGenerationToken; use tokens::tokens::sqlite::tasks::TaskId; @@ -142,7 +142,7 @@ async fn inner_loop( let _result = prompt_websocket_image(PromptWebsocketImageArgs { websocket: &mut websocket, prompt: &prompt_item.prompt, - aspect_ratio: ClientMessageAspectRatio::WideThreeByTwo, // TODO: Make configurable. + aspect_ratio: prompt_item.aspect_ratio, }).await?; let images = listen_for_websocket_images(ListenForWebsocketImagesArgs {