mirror of
https://github.com/storytold/artcraft.git
synced 2026-11-12 12:53:04 +00:00
gemini flash handler
This commit is contained in:
1 parent
79f9514cd4
commit
c386a02bc9
6 files changed
+337
-1
No files matched your search
+62
@@ -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<String>,
|
||||
|
||||
// TODO: OneOf for image_urls or image_media_tokens
|
||||
/// Image media tokens to include in the editing context.
|
||||
pub image_media_tokens: Option<Vec<MediaFileToken>>,
|
||||
|
||||
/// Size of the image to generate. Default is Square.
|
||||
pub image_size: Option<Gemini25FlashEditImageImageSize>,
|
||||
|
||||
/// Number of images to generate. Default is one.
|
||||
pub num_images: Option<Gemini25FlashEditImageNumImages>,
|
||||
|
||||
/// Quality of the image to generate. Default is High.
|
||||
pub image_quality: Option<Gemini25FlashEditImageImageQuality>,
|
||||
}
|
||||
|
||||
|
||||
#[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,
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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));
|
||||
|
||||
+258
@@ -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<GptImage1EditImageRequest>,
|
||||
server_state: web::Data<Arc<ServerState>>
|
||||
) -> Result<Json<GptImage1EditImageResponse>, 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::<Vec<_>>();
|
||||
|
||||
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,
|
||||
}))
|
||||
}
|
||||
+1
@@ -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;
|
||||
@@ -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<T, B> (app: App<T>) -> App<T>
|
||||
@@ -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()))
|
||||
|
||||
Reference in new issue
Block a user