From c386a02bc9daaa3c480fb34bdd90a60a623ac9ad Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 28 Aug 2025 01:03:08 -0400 Subject: [PATCH] gemini flash handler --- .../image/edit/gemini_25_flash_edit_image.rs | 62 +++++ .../src/generate/image/edit/mod.rs | 1 + .../public/enums/src/common/model_type.rs | 11 +- .../gemini_25_flash_edit_image_handler.rs | 258 ++++++++++++++++++ .../endpoints/generate/image/edit/mod.rs | 1 + .../http_server/routes/add_generate_routes.rs | 5 + 6 files changed, 337 insertions(+), 1 deletion(-) create mode 100644 crates/schema/public/api/artcraft_api_defs/src/generate/image/edit/gemini_25_flash_edit_image.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/edit/gemini_25_flash_edit_image_handler.rs diff --git a/crates/schema/public/api/artcraft_api_defs/src/generate/image/edit/gemini_25_flash_edit_image.rs b/crates/schema/public/api/artcraft_api_defs/src/generate/image/edit/gemini_25_flash_edit_image.rs new file mode 100644 index 0000000000..589f4ed170 --- /dev/null +++ b/crates/schema/public/api/artcraft_api_defs/src/generate/image/edit/gemini_25_flash_edit_image.rs @@ -0,0 +1,62 @@ +use serde_derive::{Deserialize, Serialize}; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::media_files::MediaFileToken; +use utoipa::ToSchema; + +pub const GEMINI_25_FLASH_EDIT_IMAGE_PATH: &str = "/v1/generate/image/edit/gemini_25_flash"; + +#[derive(Serialize, Deserialize, ToSchema)] +pub struct Gemini25FlashEditImageRequest { + /// Idempotency token to prevent duplicate requests. + pub uuid_idempotency_token: String, + + /// Text prompt to generate the image from. + pub prompt: Option, + + // TODO: OneOf for image_urls or image_media_tokens + /// Image media tokens to include in the editing context. + pub image_media_tokens: Option>, + + /// Size of the image to generate. Default is Square. + pub image_size: Option, + + /// Number of images to generate. Default is one. + pub num_images: Option, + + /// Quality of the image to generate. Default is High. + pub image_quality: Option, +} + + +#[derive(Copy, Clone, Debug, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum Gemini25FlashEditImageImageSize { + Square, // 1:1, Default + Horizontal, + Vertical, +} + +#[derive(Copy, Clone, Debug, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum Gemini25FlashEditImageNumImages { + One, // Default + Two, + Three, + Four, +} + +#[derive(Copy, Clone, Debug, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum Gemini25FlashEditImageImageQuality { + Auto, + Low, + Medium, + High, // Default +} + + +#[derive(Serialize, Deserialize, ToSchema)] +pub struct Gemini25FlashEditImageResponse { + pub success: bool, + pub inference_job_token: InferenceJobToken, +} diff --git a/crates/schema/public/api/artcraft_api_defs/src/generate/image/edit/mod.rs b/crates/schema/public/api/artcraft_api_defs/src/generate/image/edit/mod.rs index e3014dfd6a..d4f7245aa6 100644 --- a/crates/schema/public/api/artcraft_api_defs/src/generate/image/edit/mod.rs +++ b/crates/schema/public/api/artcraft_api_defs/src/generate/image/edit/mod.rs @@ -1,3 +1,4 @@ pub mod flux_pro_kontext_max_edit_image; pub mod gpt_image_1_edit_image; +pub mod gemini_25_flash_edit_image; diff --git a/crates/schema/public/enums/src/common/model_type.rs b/crates/schema/public/enums/src/common/model_type.rs index 5a6a186fc4..34c1a2ba16 100644 --- a/crates/schema/public/enums/src/common/model_type.rs +++ b/crates/schema/public/enums/src/common/model_type.rs @@ -31,6 +31,8 @@ pub enum ModelType { GptImage1, #[serde(rename = "recraft_3")] Recraft3, + #[serde(rename = "gemini_25_flash")] + Gemini25Flash, /// Midjourney without distinguishing a model type or version #[serde(rename = "midjourney")] @@ -92,6 +94,7 @@ impl ModelType { Self::FluxProKontextMax => "flux_pro_kontext_max", Self::GptImage1 => "gpt_image_1", Self::Recraft3 => "recraft_3", + Self::Gemini25Flash => "gemini_25_flash", Self::Midjourney => "midjourney", Self::MidjourneyV6 => "midjourney_v6", Self::MidjourneyV6p1 => "midjourney_v6p1", @@ -126,6 +129,7 @@ impl ModelType { "flux_pro_kontext_max" => Ok(Self::FluxProKontextMax), "gpt_image_1" => Ok(Self::GptImage1), "recraft_3" => Ok(Self::Recraft3), + "gemini_25_flash" => Ok(Self::Gemini25Flash), "midjourney" => Ok(Self::Midjourney), "midjourney_v6" => Ok(Self::MidjourneyV6), "midjourney_v6p1" => Ok(Self::MidjourneyV6p1), @@ -164,6 +168,7 @@ impl ModelType { Self::FluxProKontextMax, Self::GptImage1, Self::Recraft3, + Self::Gemini25Flash, Self::Midjourney, Self::MidjourneyV6, Self::MidjourneyV6p1, @@ -207,6 +212,7 @@ mod tests { assert_serialization(ModelType::FluxProKontextMax, "flux_pro_kontext_max"); assert_serialization(ModelType::GptImage1, "gpt_image_1"); assert_serialization(ModelType::Recraft3, "recraft_3"); + assert_serialization(ModelType::Gemini25Flash, "gemini_25_flash"); assert_serialization(ModelType::Midjourney, "midjourney"); assert_serialization(ModelType::MidjourneyV6, "midjourney_v6"); assert_serialization(ModelType::MidjourneyV6p1, "midjourney_v6p1"); @@ -238,6 +244,7 @@ mod tests { assert_eq!(ModelType::FluxProKontextMax.to_str(), "flux_pro_kontext_max"); assert_eq!(ModelType::GptImage1.to_str(), "gpt_image_1"); assert_eq!(ModelType::Recraft3.to_str(), "recraft_3"); + assert_eq!(ModelType::Gemini25Flash.to_str(), "gemini_25_flash"); assert_eq!(ModelType::Midjourney.to_str(), "midjourney"); assert_eq!(ModelType::MidjourneyV6.to_str(), "midjourney_v6"); assert_eq!(ModelType::MidjourneyV6p1.to_str(), "midjourney_v6p1"); @@ -271,6 +278,7 @@ mod tests { assert_eq!(ModelType::from_str("flux_pro_kontext_max").unwrap(), ModelType::FluxProKontextMax); assert_eq!(ModelType::from_str("gpt_image_1").unwrap(), ModelType::GptImage1); assert_eq!(ModelType::from_str("recraft_3").unwrap(), ModelType::Recraft3); + assert_eq!(ModelType::from_str("gemini_25_flash").unwrap(), ModelType::Gemini25Flash); assert_eq!(ModelType::from_str("midjourney").unwrap(), ModelType::Midjourney); assert_eq!(ModelType::from_str("midjourney_v6").unwrap(), ModelType::MidjourneyV6); assert_eq!(ModelType::from_str("midjourney_v6p1").unwrap(), ModelType::MidjourneyV6p1); @@ -293,7 +301,7 @@ mod tests { #[test] fn all_variants() { let mut variants = ModelType::all_variants(); - assert_eq!(variants.len(), 24); + assert_eq!(variants.len(), 25); // Image models assert_eq!(variants.pop_first(), Some(ModelType::Flux1Dev)); assert_eq!(variants.pop_first(), Some(ModelType::Flux1Schnell)); @@ -304,6 +312,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(ModelType::FluxProKontextMax)); assert_eq!(variants.pop_first(), Some(ModelType::GptImage1)); assert_eq!(variants.pop_first(), Some(ModelType::Recraft3)); + assert_eq!(variants.pop_first(), Some(ModelType::Gemini25Flash)); assert_eq!(variants.pop_first(), Some(ModelType::Midjourney)); assert_eq!(variants.pop_first(), Some(ModelType::MidjourneyV6)); assert_eq!(variants.pop_first(), Some(ModelType::MidjourneyV6p1)); diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/edit/gemini_25_flash_edit_image_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/edit/gemini_25_flash_edit_image_handler.rs new file mode 100644 index 0000000000..a9cef0cb80 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/edit/gemini_25_flash_edit_image_handler.rs @@ -0,0 +1,258 @@ +use std::sync::Arc; + +use crate::http_server::common_responses::common_web_error::CommonWebError; +use crate::http_server::common_responses::media::media_links_builder::MediaLinksBuilder; +use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain; +use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format; +use crate::state::server_state::ServerState; +use actix_web::web::Json; +use actix_web::{web, HttpRequest}; +use artcraft_api_defs::generate::image::edit::gpt_image_1_edit_image::{GptImage1EditImageImageQuality, GptImage1EditImageImageSize, GptImage1EditImageNumImages, GptImage1EditImageRequest, GptImage1EditImageResponse}; +use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath; +use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptContextSemanticType; +use enums::by_table::prompts::prompt_type::PromptType; +use enums::common::generation_provider::GenerationProvider; +use enums::common::model_type::ModelType; +use enums::common::visibility::Visibility; +use fal_client::creds::open_ai_api_key::OpenAiApiKey; +use fal_client::requests::webhook::image::edit::enqueue_gemini_25_flash_edit_webhook::{enqueue_gemini_25_flash_edit_webhook, Gemini25FlashEditArgs, Gemini25FlashEditNumImages}; +use fal_client::requests::webhook::image::edit::enqueue_gpt_image_1_edit_image_webhook::enqueue_gpt_image_1_edit_image_webhook; +use fal_client::requests::webhook::image::edit::enqueue_gpt_image_1_edit_image_webhook::GptEditImageNumImages; +use fal_client::requests::webhook::image::edit::enqueue_gpt_image_1_edit_image_webhook::{GptEditImageByokArgs, GptEditImageQuality, GptEditImageSize}; +use http_server_common::request::get_request_ip::get_request_ip; +use log::{error, info, warn}; +use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::insert_generic_inference_job_for_fal_queue; +use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::FalCategory; +use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::InsertGenericInferenceForFalArgs; +use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; +use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::{batch_get_media_files_by_tokens, batch_get_media_files_by_tokens_with_connection}; +use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_items::{insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem}; +use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs}; +use sqlx::Acquire; +use utoipa::ToSchema; + +/// Gemini 2.5 Flash Image Editing +#[utoipa::path( + post, + tag = "Generate Images (Edit)", + path = "/v1/generate/image/edit/gemini_25_flash", + responses( + (status = 200, description = "Success", body = GptImage1EditImageResponse), + ), + params( + ("request" = GptImage1EditImageRequest, description = "Payload for Request"), + ) +)] +pub async fn gemini_25_flash_edit_image_handler( + http_request: HttpRequest, + request: Json, + server_state: web::Data> +) -> Result, CommonWebError> { + let mut mysql_connection = server_state.mysql_pool + .acquire() + .await?; + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session_from_connection(&http_request, &mut mysql_connection) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + CommonWebError::ServerError + })?; + + let maybe_avt_token = server_state + .avt_cookie_manager + .get_avt_token_from_request(&http_request); + + // TODO: Limit usage for new accounts. Billing, free credits metering, etc. + + //let user_session = match maybe_user_session { + // Some(session) => session, + // None => { + // warn!("not logged in"); + // return Err(CommonWebError::NotAuthorized); + // } + //}; + + if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) { + return Err(CommonWebError::BadInputWithSimpleMessage(reason)); + } + + const CAN_SEE_DELETED: bool = false; + + let tokens = match request.image_media_tokens.as_ref() { + Some(tokens) => tokens, + None => { + warn!("No image media tokens provided"); + return Err(CommonWebError::BadInputWithSimpleMessage("No image media tokens provided".to_string())); + } + }; + + let result = batch_get_media_files_by_tokens_with_connection( + &mut mysql_connection, + tokens, + CAN_SEE_DELETED, + ).await; + + let media_files = match result { + Ok(files) => files, + Err(err) => { + error!("Error getting media files by tokens: {:?}", err); + return Err(CommonWebError::ServerError); + } + }; + + if media_files.len() != tokens.len() { + warn!("Wrong number of media files returned for tokens: {} found for {} tokens", media_files.len(), tokens.len()); + return Err(CommonWebError::BadInputWithSimpleMessage( + format!("Not all media files could be found. Media files found: {}, tokens provided: {}", + media_files.len(), tokens.len()))); + } + + let media_domain = get_media_domain(&http_request); + + let image_urls = media_files.iter() + .map(|file| { + let public_bucket_path = MediaFileBucketPath::from_object_hash( + &file.public_bucket_directory_hash, + file.maybe_public_bucket_prefix.as_deref(), + file.maybe_public_bucket_extension.as_deref()); + + let media_links = MediaLinksBuilder::from_media_path_and_env( + media_domain, + server_state.server_environment, + &public_bucket_path); + + media_links.cdn_url.to_string() + }) + .collect::>(); + + insert_idempotency_token(&request.uuid_idempotency_token, &mut *mysql_connection) + .await + .map_err(|err| { + error!("Error inserting idempotency token: {:?}", err); + CommonWebError::BadInputWithSimpleMessage("repeated idempotency token".to_string()) + })?; + + info!("Fal webhook URL: {}", server_state.fal.webhook_url); + + let num_images = match request.num_images { + Some(GptImage1EditImageNumImages::One) => Gemini25FlashEditNumImages::One, + Some(GptImage1EditImageNumImages::Two) => Gemini25FlashEditNumImages::Two, + Some(GptImage1EditImageNumImages::Three) => Gemini25FlashEditNumImages::Three, + Some(GptImage1EditImageNumImages::Four) => Gemini25FlashEditNumImages::Four, + None => Gemini25FlashEditNumImages::One, // Default to One + }; + + let args = Gemini25FlashEditArgs { + image_urls, + prompt: request.prompt.as_deref().unwrap_or(""), + webhook_url: &server_state.fal.webhook_url, + api_key: &server_state.fal.api_key, + num_images, + }; + + let fal_result = enqueue_gemini_25_flash_edit_webhook(args) + .await + .map_err(|err| { + warn!("Error calling enqueue_gpt_image_1_edit_image_webhook: {:?}", err); + CommonWebError::ServerError + })?; + + let external_job_id = fal_result.request_id + .ok_or_else(|| { + warn!("Fal request_id is None"); + CommonWebError::ServerError + })?; + + info!("Fal request_id: {}", external_job_id); + + let ip_address = get_request_ip(&http_request); + + let mut transaction = mysql_connection + .begin() + .await + .map_err(|err| { + error!("Error starting MySQL transaction: {:?}", err); + CommonWebError::ServerError + })?; + + // NB: Don't fail the job if the query fails. + let prompt_result = insert_prompt(InsertPromptArgs { + maybe_apriori_prompt_token: None, + prompt_type: PromptType::ArtcraftApp, + maybe_creator_user_token: maybe_user_session + .as_ref() + .map(|s| &s.user_token), + maybe_model_type: Some(ModelType::Gemini25Flash), + maybe_generation_provider: Some(GenerationProvider::Artcraft), + maybe_positive_prompt: request.prompt.as_deref(), + maybe_negative_prompt: None, + maybe_other_args: None, + creator_ip_address: &ip_address, + mysql_executor: &mut *transaction, + phantom: Default::default(), + }).await; + + let prompt_token = match prompt_result { + Ok(token) => Some(token), + Err(err) => { + warn!("Error inserting prompt: {:?}", err); + None // Don't fail the job if the prompt insertion fails. + } + }; + + if let Some(token) = prompt_token.as_ref() { + let result = insert_batch_prompt_context_items(InsertBatchArgs { + prompt_token: token.clone(), + items: media_files.iter().map(|media_file| { + PromptContextItem { + media_token: media_file.token.clone(), + context_semantic_type: PromptContextSemanticType::Imgref, + } + }).collect(), + transaction: &mut transaction, + }).await; + + if let Err(err) = result { + // NB: Fail open. + warn!("Error inserting batch prompt context items: {:?}", err); + } + } + + let db_result = insert_generic_inference_job_for_fal_queue(InsertGenericInferenceForFalArgs { + uuid_idempotency_token: &request.uuid_idempotency_token, + maybe_external_third_party_id: &external_job_id, + fal_category: FalCategory::ImageGeneration, + maybe_inference_args: None, + maybe_prompt_token: prompt_token.as_ref(), + maybe_creator_user_token: maybe_user_session.as_ref().map(|s| &s.user_token), + maybe_avt_token: maybe_avt_token.as_ref(), + creator_ip_address: &ip_address, + creator_set_visibility: Visibility::Public, + mysql_executor: &mut *transaction, + phantom: Default::default(), + }).await; + + let job_token = match db_result { + Ok(token) => token, + Err(err) => { + warn!("Error inserting generic inference job for FAL queue: {:?}", err); + return Err(CommonWebError::ServerError); + } + }; + + let _r = transaction + .commit() + .await + .map_err(|err| { + error!("Error committing MySQL transaction: {:?}", err); + CommonWebError::ServerError + })?; + + Ok(Json(GptImage1EditImageResponse { + success: true, + inference_job_token: job_token, + })) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/edit/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/edit/mod.rs index 30ca60aae0..20313b49e7 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/edit/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/edit/mod.rs @@ -1,2 +1,3 @@ pub mod flux_pro_kontext_max_edit_image_handler; +pub mod gemini_25_flash_edit_image_handler; pub mod gpt_image_1_edit_image_handler; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/routes/add_generate_routes.rs b/crates/service/web/storyteller_web/src/http_server/routes/add_generate_routes.rs index 2132e2f92f..997842b467 100644 --- a/crates/service/web/storyteller_web/src/http_server/routes/add_generate_routes.rs +++ b/crates/service/web/storyteller_web/src/http_server/routes/add_generate_routes.rs @@ -18,6 +18,7 @@ use actix_http::body::MessageBody; use actix_service::ServiceFactory; use actix_web::dev::{ServiceRequest, ServiceResponse}; use actix_web::{web, App, Error, HttpResponse}; +use crate::http_server::endpoints::generate::image::edit::gemini_25_flash_edit_image_handler::gemini_25_flash_edit_image_handler; use crate::http_server::endpoints::generate::image::inpaint::flux_dev_juggernaut_inpaint_handler::flux_dev_juggernaut_inpaint_image_handler; pub fn add_generate_routes (app: App) -> App @@ -38,6 +39,10 @@ where .route(web::post().to(flux_pro_kontext_max_edit_image_handler)) .route(web::head().to(|| HttpResponse::Ok())) ) + .service(web::resource("/gemini_25_flash") + .route(web::post().to(gemini_25_flash_edit_image_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) .service(web::resource("/gpt_image_1") .route(web::post().to(gpt_image_1_edit_image_handler)) .route(web::head().to(|| HttpResponse::Ok()))