mirror of
https://github.com/storytold/artcraft.git
synced 2026-10-29 04:52:51 -04:00
Compare commits
12
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2aa1e762d5 | ||
|
|
28a0725f34 | ||
|
|
b68d24040f | ||
|
|
3792cf82f2 | ||
|
|
f73add81ef | ||
|
|
7465da6092 | ||
|
|
fc5fda42cf | ||
|
|
0eb4ceed70 | ||
|
|
6f1045e58a | ||
|
|
2b97694ebf | ||
|
|
0f6998cb6b | ||
|
|
8076d1b830 |
@@ -51,6 +51,10 @@ then the story, then the fine print.
|
||||
- Group imports: std, external crates, internal workspace crates, `crate::` imports
|
||||
- Use fully qualified paths for one-off references; `use` for repeated references
|
||||
- Prefer specific imports over wildcards, except in `api_doc.rs` and test modules
|
||||
- When two crates export a type with the same name, alias with a suffix describing the source crate:
|
||||
`use enums::...::CommonResolution as CommonResolutionEnum;`
|
||||
`use artcraft_router::...::CommonResolution as CommonResolutionRouter;`
|
||||
This avoids long fully-qualified paths in function signatures and bodies.
|
||||
|
||||
## TypeScript / Frontend
|
||||
|
||||
|
||||
@@ -61,6 +61,7 @@ artcraft/
|
||||
- SQLx for MySQL and SQLite
|
||||
- A mix of wreq and reqwest for Rust HTTP clients
|
||||
- Never use `println!` or `eprintln!` outside of tests; use `log` crate macros instead
|
||||
- When two crates export the same type name, alias imports with a suffix: `use foo::Bar as BarFoo;`
|
||||
- TypeScript with Nx, React, Vite, Zustand, and Three.js
|
||||
- Use two spaces for indentation
|
||||
|
||||
|
||||
Generated
+2300
-2468
File diff suppressed because it is too large
Load Diff
+22
@@ -0,0 +1,22 @@
|
||||
$kind: http-request
|
||||
name: Add Wallet Banked Balance
|
||||
url: "{{api_host}}/v1/moderation/wallet/wallet_2pz2h3m2bdfvna7m3b92ccwwr/add_ba\
|
||||
nked_balance"
|
||||
method: POST
|
||||
headers:
|
||||
- key: Accept
|
||||
value: application/json
|
||||
- key: Session
|
||||
value: ""
|
||||
disabled: true
|
||||
- key: Session
|
||||
value: ""
|
||||
disabled: true
|
||||
body:
|
||||
type: json
|
||||
content: |
|
||||
{
|
||||
"title": "test alert from thread 2",
|
||||
"description": "testing testing test"
|
||||
}
|
||||
order: 3365370184335612
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
$kind: http-request
|
||||
name: Create Wallet for User
|
||||
url: "{{api_host}}/v1/moderation/wallet/create_for_user"
|
||||
method: POST
|
||||
headers:
|
||||
- key: Accept
|
||||
value: application/json
|
||||
- key: Session
|
||||
value: ""
|
||||
disabled: true
|
||||
- key: Session
|
||||
value: ""
|
||||
disabled: true
|
||||
body:
|
||||
type: json
|
||||
content: |
|
||||
{
|
||||
"user_token": "user_jg7ng4kdt2x5y", // newuser946
|
||||
"payments_namespace": "artcraft"
|
||||
}
|
||||
order: 4100053395280358
|
||||
+2
-2
@@ -89,7 +89,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_serialization_text_to_image() {
|
||||
let request = EstimateImageCostRequest {
|
||||
model: CommonImageModel::NanaBananaPro,
|
||||
model: CommonImageModel::NanoBananaPro,
|
||||
provider: GenerationProvider::Artcraft,
|
||||
generation_mode: GenerationMode::TextToImage,
|
||||
aspect_ratio: None,
|
||||
@@ -103,7 +103,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_serialization_image_edit() {
|
||||
let request = EstimateImageCostRequest {
|
||||
model: CommonImageModel::NanaBananaPro,
|
||||
model: CommonImageModel::NanoBananaPro,
|
||||
provider: GenerationProvider::Artcraft,
|
||||
generation_mode: GenerationMode::ImageEdit { count: 2 },
|
||||
aspect_ratio: None,
|
||||
|
||||
+1
-1
@@ -124,7 +124,7 @@ pub struct OmniGenImageModelDetails {
|
||||
impl Default for OmniGenImageModelDetails {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
model: CommonImageModel::NanaBananaPro,
|
||||
model: CommonImageModel::NanoBananaPro,
|
||||
full_name: None,
|
||||
text_prompt_supported: None,
|
||||
text_prompt_max_length: None,
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
|
||||
#[derive(Copy, Clone)]
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub enum AudioListRef<'a> {
|
||||
MediaFileTokens(&'a Vec<MediaFileToken>),
|
||||
Urls(&'a Vec<String>),
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use tokens::tokens::characters::CharacterToken;
|
||||
|
||||
/// Forward-compatible character reference list.
|
||||
#[derive(Copy, Clone)]
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub enum CharacterListRef<'a> {
|
||||
CharacterTokens(&'a Vec<CharacterToken>),
|
||||
|
||||
|
||||
@@ -16,11 +16,11 @@ pub enum CommonImageModel {
|
||||
#[serde(rename = "gpt_image_1p5")]
|
||||
GptImage1p5,
|
||||
#[serde(rename = "nano_banana")]
|
||||
NanaBanana,
|
||||
NanoBanana,
|
||||
#[serde(rename = "nano_banana_2")]
|
||||
NanaBanana2,
|
||||
NanoBanana2,
|
||||
#[serde(rename = "nano_banana_pro")]
|
||||
NanaBananaPro,
|
||||
NanoBananaPro,
|
||||
#[serde(rename = "seedream_4")]
|
||||
Seedream4,
|
||||
#[serde(rename = "seedream_4p5")]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
|
||||
#[derive(Copy, Clone)]
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub enum ImageListRef<'a> {
|
||||
MediaFileTokens(&'a Vec<MediaFileToken>),
|
||||
Urls(&'a Vec<String>),
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
|
||||
#[derive(Copy, Clone)]
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub enum ImageRef<'a> {
|
||||
MediaFileToken(&'a MediaFileToken),
|
||||
Url(&'a str),
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
|
||||
#[derive(Copy, Clone)]
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub enum VideoListRef<'a> {
|
||||
MediaFileTokens(&'a Vec<MediaFileToken>),
|
||||
Urls(&'a Vec<String>),
|
||||
|
||||
+1
-1
@@ -37,7 +37,7 @@ mod tests {
|
||||
|
||||
fn estimate_usd_cents(image_batch_count: u16) -> u64 {
|
||||
let request = GenerateImageRequest {
|
||||
model: CommonImageModel::NanaBanana,
|
||||
model: CommonImageModel::NanoBanana,
|
||||
provider: Provider::Artcraft,
|
||||
prompt: None,
|
||||
image_inputs: None,
|
||||
|
||||
+1
-1
@@ -45,7 +45,7 @@ mod tests {
|
||||
|
||||
fn estimate_usd_cents(resolution: Option<CommonResolution>, image_batch_count: u16) -> u64 {
|
||||
let request = GenerateImageRequest {
|
||||
model: CommonImageModel::NanaBanana2,
|
||||
model: CommonImageModel::NanoBanana2,
|
||||
provider: Provider::Artcraft,
|
||||
prompt: None,
|
||||
image_inputs: None,
|
||||
|
||||
+1
-1
@@ -44,7 +44,7 @@ mod tests {
|
||||
|
||||
fn estimate_usd_cents(resolution: Option<CommonResolution>, image_batch_count: u16) -> u64 {
|
||||
let request = GenerateImageRequest {
|
||||
model: CommonImageModel::NanaBananaPro,
|
||||
model: CommonImageModel::NanoBananaPro,
|
||||
provider: Provider::Artcraft,
|
||||
prompt: None,
|
||||
image_inputs: None,
|
||||
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
use crate::generate::generate_image::image_generation_cost_estimate::ImageGenerationCostEstimate;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_flux_1_dev::PlanFalFlux1Dev;
|
||||
|
||||
pub(crate) fn estimate_image_cost_fal_flux_1_dev(
|
||||
plan: &PlanFalFlux1Dev<'_>,
|
||||
) -> ImageGenerationCostEstimate {
|
||||
// Pricing: $0.025 per megapixel, billed by rounding up to nearest megapixel.
|
||||
// Default image sizes are ~1MP (1024x1024), so ~$0.025/image = 3 cents rounded up.
|
||||
//
|
||||
// Fal docs:
|
||||
// Your request will cost $0.025 per megapixel, with images billed by rounding up
|
||||
// to the nearest megapixel.
|
||||
//
|
||||
let cost_per_image: u64 = 3;
|
||||
let num_images = plan.num_images.as_u64();
|
||||
let cost_in_usd_cents = cost_per_image * num_images;
|
||||
|
||||
ImageGenerationCostEstimate {
|
||||
cost_in_credits: Some(cost_in_usd_cents),
|
||||
cost_in_usd_cents: Some(cost_in_usd_cents),
|
||||
is_free: false,
|
||||
is_unlimited: false,
|
||||
is_rate_limited: false,
|
||||
has_watermark: false,
|
||||
}
|
||||
}
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
use crate::generate::generate_image::image_generation_cost_estimate::ImageGenerationCostEstimate;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_flux_1_schnell::PlanFalFlux1Schnell;
|
||||
|
||||
pub(crate) fn estimate_image_cost_fal_flux_1_schnell(
|
||||
plan: &PlanFalFlux1Schnell<'_>,
|
||||
) -> ImageGenerationCostEstimate {
|
||||
// Pricing: $0.003 per megapixel, billed by rounding up to nearest megapixel.
|
||||
// Default image sizes are ~1MP (1024x1024), so ~$0.003/image = 1 cent rounded up.
|
||||
//
|
||||
// Fal docs:
|
||||
// Your request will cost $0.003 per megapixel, with images billed by rounding up
|
||||
// to the nearest megapixel.
|
||||
//
|
||||
let cost_per_image: u64 = 1;
|
||||
let num_images = plan.num_images.as_u64();
|
||||
let cost_in_usd_cents = cost_per_image * num_images;
|
||||
|
||||
ImageGenerationCostEstimate {
|
||||
cost_in_credits: Some(cost_in_usd_cents),
|
||||
cost_in_usd_cents: Some(cost_in_usd_cents),
|
||||
is_free: false,
|
||||
is_unlimited: false,
|
||||
is_rate_limited: false,
|
||||
has_watermark: false,
|
||||
}
|
||||
}
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
use crate::generate::generate_image::image_generation_cost_estimate::ImageGenerationCostEstimate;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana::{
|
||||
FalNbNumImages, PlanFalNanoBanana,
|
||||
};
|
||||
|
||||
pub(crate) fn estimate_image_cost_fal_nano_banana(
|
||||
plan: &PlanFalNanoBanana<'_>,
|
||||
) -> ImageGenerationCostEstimate {
|
||||
// Pricing: $0.039/image (no resolution multiplier — Gemini 25 Flash has no resolution option).
|
||||
// Rounded to 4 cents per image.
|
||||
//
|
||||
// Fal docs:
|
||||
// Your request will cost $0.039 per image. For $1.00, you can run this model 25 times.
|
||||
//
|
||||
let cost_per_image: u64 = 4;
|
||||
|
||||
let num_images: u64 = match plan.num_images {
|
||||
FalNbNumImages::One => 1,
|
||||
FalNbNumImages::Two => 2,
|
||||
FalNbNumImages::Three => 3,
|
||||
FalNbNumImages::Four => 4,
|
||||
};
|
||||
|
||||
let cost_in_usd_cents = cost_per_image * num_images;
|
||||
|
||||
ImageGenerationCostEstimate {
|
||||
cost_in_credits: Some(cost_in_usd_cents),
|
||||
cost_in_usd_cents: Some(cost_in_usd_cents),
|
||||
is_free: false,
|
||||
is_unlimited: false,
|
||||
is_rate_limited: false,
|
||||
has_watermark: false,
|
||||
}
|
||||
}
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
use crate::generate::generate_image::image_generation_cost_estimate::ImageGenerationCostEstimate;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana_2::{
|
||||
FalNb2NumImages, FalNb2Resolution, PlanFalNanaBanana2,
|
||||
};
|
||||
|
||||
pub(crate) fn estimate_image_cost_fal_nano_banana_2(
|
||||
plan: &PlanFalNanaBanana2<'_>,
|
||||
) -> ImageGenerationCostEstimate {
|
||||
// Pricing: $0.08/image base (1K). Multipliers: 0.5K = 0.75x, 2K = 1.5x, 4K = 2x.
|
||||
// (Stored as USD cents.)
|
||||
//
|
||||
// Fal docs:
|
||||
// Your request will cost $0.08 per image. For $1.00, you can run this model 12 times.
|
||||
// 2K and 4K outputs will be charged at 1.5 times and 2 times the standard rate, respectively.
|
||||
// 0.5K (512px) resolution outputs will be charged at 0.75 times the standard rate. If web
|
||||
// search is used, an additional $0.015 will be charged. If high thinking is used,
|
||||
// an additional $0.002 will be charged. Note: Pricing is subject to change.
|
||||
//
|
||||
let cost_per_image: u64 = match plan.resolution {
|
||||
Some(FalNb2Resolution::HalfK) => 6, // $0.06 = $0.08 * 0.75
|
||||
Some(FalNb2Resolution::OneK) => 8, // $0.08 base
|
||||
Some(FalNb2Resolution::TwoK) => 12, // $0.12 = $0.08 * 1.5
|
||||
Some(FalNb2Resolution::FourK) => 16, // $0.16 = $0.08 * 2.0
|
||||
None => 8, // default = 1K base
|
||||
};
|
||||
|
||||
let num_images: u64 = match plan.num_images {
|
||||
FalNb2NumImages::One => 1,
|
||||
FalNb2NumImages::Two => 2,
|
||||
FalNb2NumImages::Three => 3,
|
||||
FalNb2NumImages::Four => 4,
|
||||
};
|
||||
|
||||
let cost_in_usd_cents = cost_per_image * num_images;
|
||||
|
||||
ImageGenerationCostEstimate {
|
||||
cost_in_credits: Some(cost_in_usd_cents),
|
||||
cost_in_usd_cents: Some(cost_in_usd_cents),
|
||||
is_free: false,
|
||||
is_unlimited: false,
|
||||
is_rate_limited: false,
|
||||
has_watermark: false,
|
||||
}
|
||||
}
|
||||
+1
-1
@@ -42,7 +42,7 @@ mod tests {
|
||||
|
||||
fn estimate_usd_cents(resolution: Option<CommonResolution>, image_batch_count: u16) -> u64 {
|
||||
let request = GenerateImageRequest {
|
||||
model: CommonImageModel::NanaBananaPro,
|
||||
model: CommonImageModel::NanoBananaPro,
|
||||
provider: Provider::Fal,
|
||||
prompt: None,
|
||||
image_inputs: None,
|
||||
|
||||
+4
@@ -1 +1,5 @@
|
||||
pub(crate) mod estimate_image_cost_fal_flux_1_dev;
|
||||
pub(crate) mod estimate_image_cost_fal_flux_1_schnell;
|
||||
pub(crate) mod estimate_image_cost_fal_nano_banana;
|
||||
pub(crate) mod estimate_image_cost_fal_nano_banana_2;
|
||||
pub(crate) mod estimate_image_cost_fal_nano_banana_pro;
|
||||
|
||||
+49
@@ -0,0 +1,49 @@
|
||||
use crate::client::router_fal_client::RouterFalClient;
|
||||
use crate::errors::artcraft_router_error::ArtcraftRouterError;
|
||||
use crate::errors::provider_error::ProviderError;
|
||||
use crate::generate::generate_image::generate_image_response::{
|
||||
FalImageResponsePayload, GenerateImageResponse,
|
||||
};
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_flux_1_dev::PlanFalFlux1Dev;
|
||||
use fal_client::requests::webhook::image::edit::enqueue_flux_1_dev_edit_image_webhook::{
|
||||
enqueue_flux_1_dev_edit_image_webhook, Flux1DevEditImageArgs,
|
||||
};
|
||||
use fal_client::requests::webhook::image::text::enqueue_flux_1_dev_text_to_image_webhook::{
|
||||
enqueue_flux_1_dev_text_to_image_webhook, Flux1DevArgs,
|
||||
};
|
||||
|
||||
pub async fn execute_fal_flux_1_dev(
|
||||
plan: &PlanFalFlux1Dev<'_>,
|
||||
fal_client: &RouterFalClient,
|
||||
) -> Result<GenerateImageResponse, ArtcraftRouterError> {
|
||||
let webhook_response = if let Some(image_url) = &plan.maybe_image_url {
|
||||
// Image-to-image mode
|
||||
let args = Flux1DevEditImageArgs {
|
||||
prompt: plan.prompt.unwrap_or(""),
|
||||
image_url: image_url.clone(),
|
||||
num_images: plan.num_images.to_edit(),
|
||||
webhook_url: fal_client.webhook_url.as_str(),
|
||||
api_key: &fal_client.api_key,
|
||||
};
|
||||
enqueue_flux_1_dev_edit_image_webhook(args)
|
||||
.await
|
||||
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?
|
||||
} else {
|
||||
// Text-to-image mode
|
||||
let args = Flux1DevArgs {
|
||||
prompt: plan.prompt.unwrap_or(""),
|
||||
aspect_ratio: plan.aspect_ratio,
|
||||
num_images: plan.num_images.to_t2i(),
|
||||
webhook_url: fal_client.webhook_url.as_str(),
|
||||
api_key: &fal_client.api_key,
|
||||
};
|
||||
enqueue_flux_1_dev_text_to_image_webhook(args)
|
||||
.await
|
||||
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?
|
||||
};
|
||||
|
||||
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
|
||||
request_id: webhook_response.request_id,
|
||||
gateway_request_id: webhook_response.gateway_request_id,
|
||||
}))
|
||||
}
|
||||
+49
@@ -0,0 +1,49 @@
|
||||
use crate::client::router_fal_client::RouterFalClient;
|
||||
use crate::errors::artcraft_router_error::ArtcraftRouterError;
|
||||
use crate::errors::provider_error::ProviderError;
|
||||
use crate::generate::generate_image::generate_image_response::{
|
||||
FalImageResponsePayload, GenerateImageResponse,
|
||||
};
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_flux_1_schnell::PlanFalFlux1Schnell;
|
||||
use fal_client::requests::webhook::image::edit::enqueue_flux_1_schnell_edit_image_webhook::{
|
||||
enqueue_flux_1_schnell_edit_image_webhook, Flux1SchnellEditImageArgs,
|
||||
};
|
||||
use fal_client::requests::webhook::image::text::enqueue_flux_1_schnell_text_to_image_webhook::{
|
||||
enqueue_flux_1_schnell_text_to_image_webhook, Flux1SchnellArgs,
|
||||
};
|
||||
|
||||
pub async fn execute_fal_flux_1_schnell(
|
||||
plan: &PlanFalFlux1Schnell<'_>,
|
||||
fal_client: &RouterFalClient,
|
||||
) -> Result<GenerateImageResponse, ArtcraftRouterError> {
|
||||
let webhook_response = if let Some(image_url) = &plan.maybe_image_url {
|
||||
// Image-to-image (redux) mode
|
||||
let args = Flux1SchnellEditImageArgs {
|
||||
image_url: image_url.clone(),
|
||||
num_images: plan.num_images.to_edit(),
|
||||
image_size: plan.edit_image_size,
|
||||
webhook_url: fal_client.webhook_url.as_str(),
|
||||
api_key: &fal_client.api_key,
|
||||
};
|
||||
enqueue_flux_1_schnell_edit_image_webhook(args)
|
||||
.await
|
||||
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?
|
||||
} else {
|
||||
// Text-to-image mode
|
||||
let args = Flux1SchnellArgs {
|
||||
prompt: plan.prompt.unwrap_or(""),
|
||||
aspect_ratio: plan.t2i_aspect_ratio,
|
||||
num_images: plan.num_images.to_t2i(),
|
||||
webhook_url: fal_client.webhook_url.as_str(),
|
||||
api_key: &fal_client.api_key,
|
||||
};
|
||||
enqueue_flux_1_schnell_text_to_image_webhook(args)
|
||||
.await
|
||||
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?
|
||||
};
|
||||
|
||||
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
|
||||
request_id: webhook_response.request_id,
|
||||
gateway_request_id: webhook_response.gateway_request_id,
|
||||
}))
|
||||
}
|
||||
+72
@@ -0,0 +1,72 @@
|
||||
use crate::client::router_fal_client::RouterFalClient;
|
||||
use crate::errors::artcraft_router_error::ArtcraftRouterError;
|
||||
use crate::errors::provider_error::ProviderError;
|
||||
use crate::generate::generate_image::generate_image_response::{
|
||||
FalImageResponsePayload, GenerateImageResponse,
|
||||
};
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana::{
|
||||
FalNbNumImages, PlanFalNanoBanana,
|
||||
};
|
||||
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::text::enqueue_gemini_25_flash_text_to_image_webhook::{
|
||||
enqueue_gemini_25_flash_text_to_image_webhook, Gemini25FlashTextToImageArgs,
|
||||
Gemini25FlashTextToImageNumImages,
|
||||
};
|
||||
|
||||
pub async fn execute_fal_nano_banana(
|
||||
plan: &PlanFalNanoBanana<'_>,
|
||||
fal_client: &RouterFalClient,
|
||||
) -> Result<GenerateImageResponse, ArtcraftRouterError> {
|
||||
let webhook_response = if plan.image_urls.is_empty() {
|
||||
// Text-to-image mode
|
||||
let args = Gemini25FlashTextToImageArgs {
|
||||
prompt: plan.prompt.unwrap_or(""),
|
||||
num_images: to_t2i_num_images(plan.num_images),
|
||||
aspect_ratio: plan.t2i_aspect_ratio,
|
||||
webhook_url: fal_client.webhook_url.as_str(),
|
||||
api_key: &fal_client.api_key,
|
||||
};
|
||||
enqueue_gemini_25_flash_text_to_image_webhook(args)
|
||||
.await
|
||||
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?
|
||||
} else {
|
||||
// Image-edit mode
|
||||
let args = Gemini25FlashEditArgs {
|
||||
prompt: plan.prompt.unwrap_or(""),
|
||||
image_urls: plan.image_urls.clone(),
|
||||
num_images: to_edit_num_images(plan.num_images),
|
||||
aspect_ratio: plan.edit_aspect_ratio,
|
||||
webhook_url: fal_client.webhook_url.as_str(),
|
||||
api_key: &fal_client.api_key,
|
||||
};
|
||||
enqueue_gemini_25_flash_edit_webhook(args)
|
||||
.await
|
||||
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?
|
||||
};
|
||||
|
||||
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
|
||||
request_id: webhook_response.request_id,
|
||||
gateway_request_id: webhook_response.gateway_request_id,
|
||||
}))
|
||||
}
|
||||
|
||||
fn to_t2i_num_images(n: FalNbNumImages) -> Gemini25FlashTextToImageNumImages {
|
||||
match n {
|
||||
FalNbNumImages::One => Gemini25FlashTextToImageNumImages::One,
|
||||
FalNbNumImages::Two => Gemini25FlashTextToImageNumImages::Two,
|
||||
FalNbNumImages::Three => Gemini25FlashTextToImageNumImages::Three,
|
||||
FalNbNumImages::Four => Gemini25FlashTextToImageNumImages::Four,
|
||||
}
|
||||
}
|
||||
|
||||
fn to_edit_num_images(n: FalNbNumImages) -> Gemini25FlashEditNumImages {
|
||||
match n {
|
||||
FalNbNumImages::One => Gemini25FlashEditNumImages::One,
|
||||
FalNbNumImages::Two => Gemini25FlashEditNumImages::Two,
|
||||
FalNbNumImages::Three => Gemini25FlashEditNumImages::Three,
|
||||
FalNbNumImages::Four => Gemini25FlashEditNumImages::Four,
|
||||
}
|
||||
}
|
||||
+92
@@ -0,0 +1,92 @@
|
||||
use crate::client::router_fal_client::RouterFalClient;
|
||||
use crate::errors::artcraft_router_error::ArtcraftRouterError;
|
||||
use crate::errors::provider_error::ProviderError;
|
||||
use crate::generate::generate_image::generate_image_response::{
|
||||
FalImageResponsePayload, GenerateImageResponse,
|
||||
};
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana_2::{
|
||||
FalNb2NumImages, FalNb2Resolution, PlanFalNanaBanana2,
|
||||
};
|
||||
use fal_client::requests::webhook::image::edit::enqueue_nano_banana_2_edit_image_webhook::{
|
||||
enqueue_nano_banana_2_edit_image_webhook, EnqueueNanoBanana2EditImageArgs,
|
||||
EnqueueNanoBanana2EditImageNumImages, EnqueueNanoBanana2EditImageResolution,
|
||||
};
|
||||
use fal_client::requests::webhook::image::text::enqueue_nano_banana_2_text_to_image_webhook::{
|
||||
enqueue_nano_banana_2_text_to_image_webhook, EnqueueNanoBanana2TextToImageArgs,
|
||||
EnqueueNanoBanana2TextToImageNumImages, EnqueueNanoBanana2TextToImageResolution,
|
||||
};
|
||||
|
||||
pub async fn execute_fal_nano_banana_2(
|
||||
plan: &PlanFalNanaBanana2<'_>,
|
||||
fal_client: &RouterFalClient,
|
||||
) -> Result<GenerateImageResponse, ArtcraftRouterError> {
|
||||
let webhook_response = if plan.image_urls.is_empty() {
|
||||
// Text-to-image mode
|
||||
let args = EnqueueNanoBanana2TextToImageArgs {
|
||||
prompt: plan.prompt.unwrap_or(""),
|
||||
num_images: to_t2i_num_images(plan.num_images),
|
||||
resolution: plan.resolution.map(to_t2i_resolution),
|
||||
aspect_ratio: plan.t2i_aspect_ratio,
|
||||
webhook_url: fal_client.webhook_url.as_str(),
|
||||
api_key: &fal_client.api_key,
|
||||
};
|
||||
enqueue_nano_banana_2_text_to_image_webhook(args)
|
||||
.await
|
||||
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?
|
||||
} else {
|
||||
// Image-edit mode
|
||||
let args = EnqueueNanoBanana2EditImageArgs {
|
||||
prompt: plan.prompt.unwrap_or(""),
|
||||
image_urls: plan.image_urls.clone(),
|
||||
num_images: to_edit_num_images(plan.num_images),
|
||||
resolution: plan.resolution.map(to_edit_resolution),
|
||||
aspect_ratio: plan.edit_aspect_ratio,
|
||||
webhook_url: fal_client.webhook_url.as_str(),
|
||||
api_key: &fal_client.api_key,
|
||||
};
|
||||
enqueue_nano_banana_2_edit_image_webhook(args)
|
||||
.await
|
||||
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?
|
||||
};
|
||||
|
||||
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
|
||||
request_id: webhook_response.request_id,
|
||||
gateway_request_id: webhook_response.gateway_request_id,
|
||||
}))
|
||||
}
|
||||
|
||||
fn to_t2i_num_images(n: FalNb2NumImages) -> EnqueueNanoBanana2TextToImageNumImages {
|
||||
match n {
|
||||
FalNb2NumImages::One => EnqueueNanoBanana2TextToImageNumImages::One,
|
||||
FalNb2NumImages::Two => EnqueueNanoBanana2TextToImageNumImages::Two,
|
||||
FalNb2NumImages::Three => EnqueueNanoBanana2TextToImageNumImages::Three,
|
||||
FalNb2NumImages::Four => EnqueueNanoBanana2TextToImageNumImages::Four,
|
||||
}
|
||||
}
|
||||
|
||||
fn to_edit_num_images(n: FalNb2NumImages) -> EnqueueNanoBanana2EditImageNumImages {
|
||||
match n {
|
||||
FalNb2NumImages::One => EnqueueNanoBanana2EditImageNumImages::One,
|
||||
FalNb2NumImages::Two => EnqueueNanoBanana2EditImageNumImages::Two,
|
||||
FalNb2NumImages::Three => EnqueueNanoBanana2EditImageNumImages::Three,
|
||||
FalNb2NumImages::Four => EnqueueNanoBanana2EditImageNumImages::Four,
|
||||
}
|
||||
}
|
||||
|
||||
fn to_t2i_resolution(r: FalNb2Resolution) -> EnqueueNanoBanana2TextToImageResolution {
|
||||
match r {
|
||||
FalNb2Resolution::HalfK => EnqueueNanoBanana2TextToImageResolution::HalfK,
|
||||
FalNb2Resolution::OneK => EnqueueNanoBanana2TextToImageResolution::OneK,
|
||||
FalNb2Resolution::TwoK => EnqueueNanoBanana2TextToImageResolution::TwoK,
|
||||
FalNb2Resolution::FourK => EnqueueNanoBanana2TextToImageResolution::FourK,
|
||||
}
|
||||
}
|
||||
|
||||
fn to_edit_resolution(r: FalNb2Resolution) -> EnqueueNanoBanana2EditImageResolution {
|
||||
match r {
|
||||
FalNb2Resolution::HalfK => EnqueueNanoBanana2EditImageResolution::HalfK,
|
||||
FalNb2Resolution::OneK => EnqueueNanoBanana2EditImageResolution::OneK,
|
||||
FalNb2Resolution::TwoK => EnqueueNanoBanana2EditImageResolution::TwoK,
|
||||
FalNb2Resolution::FourK => EnqueueNanoBanana2EditImageResolution::FourK,
|
||||
}
|
||||
}
|
||||
+4
@@ -1 +1,5 @@
|
||||
pub mod generate_image_fal_flux_1_dev;
|
||||
pub mod generate_image_fal_flux_1_schnell;
|
||||
pub mod generate_image_fal_nano_banana;
|
||||
pub mod generate_image_fal_nano_banana_2;
|
||||
pub mod generate_image_fal_nano_banana_pro;
|
||||
|
||||
+15
-6
@@ -10,19 +10,24 @@ use crate::errors::client_error::ClientError;
|
||||
use crate::generate::generate_image::image_generation_plan::ImageGenerationPlan;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_flux_1_dev::plan_generate_image_artcraft_flux_1_dev;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_flux_1_schnell::plan_generate_image_artcraft_flux_1_schnell;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_flux_2_lora_angles::plan_generate_image_artcraft_flux_2_lora_angles;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_flux_pro_1p1::plan_generate_image_artcraft_flux_pro_1p1;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_flux_pro_1p1_ultra::plan_generate_image_artcraft_flux_pro_1p1_ultra;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_gpt_image_1p5::plan_generate_image_artcraft_gpt_image_1p5;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_nano_banana::plan_generate_image_artcraft_nano_banana;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_nano_banana_2::plan_generate_image_artcraft_nano_banana_2;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_nano_banana_pro::plan_generate_image_artcraft_nano_banana_pro;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_qwen_edit_2511_angles::plan_generate_image_artcraft_qwen_edit_2511_angles;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_seedream_4::plan_generate_image_artcraft_seedream_4;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_seedream_4p5::plan_generate_image_artcraft_seedream_4p5;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_seedream_5_lite::plan_generate_image_artcraft_seedream_5_lite;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_qwen_edit_2511_angles::plan_generate_image_artcraft_qwen_edit_2511_angles;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_flux_2_lora_angles::plan_generate_image_artcraft_flux_2_lora_angles;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_flux_1_dev::plan_generate_image_fal_flux_1_dev;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_flux_1_schnell::plan_generate_image_fal_flux_1_schnell;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana::plan_generate_image_fal_nano_banana;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana_2::plan_generate_image_fal_nano_banana_2;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana_pro::plan_generate_image_fal_nano_banana_pro;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct GenerateImageRequest<'a> {
|
||||
/// Which model to use.
|
||||
pub model: CommonImageModel,
|
||||
@@ -86,9 +91,9 @@ impl<'a> GenerateImageRequest<'a> {
|
||||
CommonImageModel::FluxPro11 => plan_generate_image_artcraft_flux_pro_1p1(self),
|
||||
CommonImageModel::FluxPro11Ultra => plan_generate_image_artcraft_flux_pro_1p1_ultra(self),
|
||||
CommonImageModel::GptImage1p5 => plan_generate_image_artcraft_gpt_image_1p5(self),
|
||||
CommonImageModel::NanaBanana => plan_generate_image_artcraft_nano_banana(self),
|
||||
CommonImageModel::NanaBanana2 => plan_generate_image_artcraft_nano_banana_2(self),
|
||||
CommonImageModel::NanaBananaPro => plan_generate_image_artcraft_nano_banana_pro(self),
|
||||
CommonImageModel::NanoBanana => plan_generate_image_artcraft_nano_banana(self),
|
||||
CommonImageModel::NanoBanana2 => plan_generate_image_artcraft_nano_banana_2(self),
|
||||
CommonImageModel::NanoBananaPro => plan_generate_image_artcraft_nano_banana_pro(self),
|
||||
CommonImageModel::Seedream4 => plan_generate_image_artcraft_seedream_4(self),
|
||||
CommonImageModel::Seedream4p5 => plan_generate_image_artcraft_seedream_4p5(self),
|
||||
CommonImageModel::Seedream5Lite => plan_generate_image_artcraft_seedream_5_lite(self),
|
||||
@@ -99,7 +104,11 @@ impl<'a> GenerateImageRequest<'a> {
|
||||
|
||||
fn build_fal(&self) -> Result<ImageGenerationPlan<'_>, ArtcraftRouterError> {
|
||||
match self.model {
|
||||
CommonImageModel::NanaBananaPro => plan_generate_image_fal_nano_banana_pro(self),
|
||||
CommonImageModel::Flux1Dev => plan_generate_image_fal_flux_1_dev(self),
|
||||
CommonImageModel::Flux1Schnell => plan_generate_image_fal_flux_1_schnell(self),
|
||||
CommonImageModel::NanoBanana => plan_generate_image_fal_nano_banana(self),
|
||||
CommonImageModel::NanoBanana2 => plan_generate_image_fal_nano_banana_2(self),
|
||||
CommonImageModel::NanoBananaPro => plan_generate_image_fal_nano_banana_pro(self),
|
||||
_ => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "provider",
|
||||
|
||||
+44
@@ -13,6 +13,10 @@ use crate::generate::generate_image::cost::artcraft::estimate_image_cost_artcraf
|
||||
use crate::generate::generate_image::cost::artcraft::estimate_image_cost_artcraft_seedream_5_lite::estimate_image_cost_artcraft_seedream_5_lite;
|
||||
use crate::generate::generate_image::cost::artcraft::estimate_image_cost_artcraft_qwen_edit_2511_angles::estimate_image_cost_artcraft_qwen_edit_2511_angles;
|
||||
use crate::generate::generate_image::cost::artcraft::estimate_image_cost_artcraft_flux_2_lora_angles::estimate_image_cost_artcraft_flux_2_lora_angles;
|
||||
use crate::generate::generate_image::cost::fal::estimate_image_cost_fal_flux_1_dev::estimate_image_cost_fal_flux_1_dev;
|
||||
use crate::generate::generate_image::cost::fal::estimate_image_cost_fal_flux_1_schnell::estimate_image_cost_fal_flux_1_schnell;
|
||||
use crate::generate::generate_image::cost::fal::estimate_image_cost_fal_nano_banana::estimate_image_cost_fal_nano_banana;
|
||||
use crate::generate::generate_image::cost::fal::estimate_image_cost_fal_nano_banana_2::estimate_image_cost_fal_nano_banana_2;
|
||||
use crate::generate::generate_image::cost::fal::estimate_image_cost_fal_nano_banana_pro::estimate_image_cost_fal_nano_banana_pro;
|
||||
use crate::generate::generate_image::execute::artcraft::generate_image_artcraft_flux_1_dev::execute_artcraft_flux_1_dev;
|
||||
use crate::generate::generate_image::execute::artcraft::generate_image_artcraft_flux_1_schnell::execute_artcraft_flux_1_schnell;
|
||||
@@ -27,6 +31,10 @@ use crate::generate::generate_image::execute::artcraft::generate_image_artcraft_
|
||||
use crate::generate::generate_image::execute::artcraft::generate_image_artcraft_seedream_5_lite::execute_artcraft_seedream_5_lite;
|
||||
use crate::generate::generate_image::execute::artcraft::generate_image_artcraft_qwen_edit_2511_angles::execute_artcraft_qwen_edit_2511_angles;
|
||||
use crate::generate::generate_image::execute::artcraft::generate_image_artcraft_flux_2_lora_angles::execute_artcraft_flux_2_lora_angles;
|
||||
use crate::generate::generate_image::execute::fal::generate_image_fal_flux_1_dev::execute_fal_flux_1_dev;
|
||||
use crate::generate::generate_image::execute::fal::generate_image_fal_flux_1_schnell::execute_fal_flux_1_schnell;
|
||||
use crate::generate::generate_image::execute::fal::generate_image_fal_nano_banana::execute_fal_nano_banana;
|
||||
use crate::generate::generate_image::execute::fal::generate_image_fal_nano_banana_2::execute_fal_nano_banana_2;
|
||||
use crate::generate::generate_image::execute::fal::generate_image_fal_nano_banana_pro::execute_fal_nano_banana_pro;
|
||||
use crate::generate::generate_image::generate_image_response::GenerateImageResponse;
|
||||
use crate::generate::generate_image::image_generation_cost_estimate::ImageGenerationCostEstimate;
|
||||
@@ -43,6 +51,10 @@ use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraf
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_seedream_5_lite::PlanArtcraftSeedream5Lite;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_qwen_edit_2511_angles::PlanArtcraftQwenEdit2511Angles;
|
||||
use crate::generate::generate_image::plan::artcraft::plan_generate_image_artcraft_flux_2_lora_angles::PlanArtcraftFlux2LoraAngles;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_flux_1_dev::PlanFalFlux1Dev;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_flux_1_schnell::PlanFalFlux1Schnell;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana::PlanFalNanoBanana;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana_2::PlanFalNanaBanana2;
|
||||
use crate::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana_pro::PlanFalNanaBananaPro;
|
||||
|
||||
#[derive(Debug)]
|
||||
@@ -60,6 +72,10 @@ pub enum ImageGenerationPlan<'a> {
|
||||
ArtcraftSeedream5Lite(PlanArtcraftSeedream5Lite<'a>),
|
||||
ArtcraftQwenEdit2511Angles(PlanArtcraftQwenEdit2511Angles<'a>),
|
||||
ArtcraftFlux2LoraAngles(PlanArtcraftFlux2LoraAngles<'a>),
|
||||
FalFlux1Dev(PlanFalFlux1Dev<'a>),
|
||||
FalFlux1Schnell(PlanFalFlux1Schnell<'a>),
|
||||
FalNanoBanana(PlanFalNanoBanana<'a>),
|
||||
FalNanaBanana2(PlanFalNanaBanana2<'a>),
|
||||
FalNanaBananaPro(PlanFalNanaBananaPro<'a>),
|
||||
}
|
||||
|
||||
@@ -121,6 +137,22 @@ impl<'a> ImageGenerationPlan<'a> {
|
||||
let artcraft_client = client.get_artcraft_client_ref()?;
|
||||
execute_artcraft_flux_2_lora_angles(plan, artcraft_client).await
|
||||
}
|
||||
ImageGenerationPlan::FalFlux1Dev(plan) => {
|
||||
let fal_client = client.get_fal_client_ref()?;
|
||||
execute_fal_flux_1_dev(plan, fal_client).await
|
||||
}
|
||||
ImageGenerationPlan::FalFlux1Schnell(plan) => {
|
||||
let fal_client = client.get_fal_client_ref()?;
|
||||
execute_fal_flux_1_schnell(plan, fal_client).await
|
||||
}
|
||||
ImageGenerationPlan::FalNanoBanana(plan) => {
|
||||
let fal_client = client.get_fal_client_ref()?;
|
||||
execute_fal_nano_banana(plan, fal_client).await
|
||||
}
|
||||
ImageGenerationPlan::FalNanaBanana2(plan) => {
|
||||
let fal_client = client.get_fal_client_ref()?;
|
||||
execute_fal_nano_banana_2(plan, fal_client).await
|
||||
}
|
||||
ImageGenerationPlan::FalNanaBananaPro(plan) => {
|
||||
let fal_client = client.get_fal_client_ref()?;
|
||||
execute_fal_nano_banana_pro(plan, fal_client).await
|
||||
@@ -169,6 +201,18 @@ impl<'a> ImageGenerationPlan<'a> {
|
||||
ImageGenerationPlan::ArtcraftFlux2LoraAngles(plan) => {
|
||||
estimate_image_cost_artcraft_flux_2_lora_angles(plan)
|
||||
}
|
||||
ImageGenerationPlan::FalFlux1Dev(plan) => {
|
||||
estimate_image_cost_fal_flux_1_dev(plan)
|
||||
}
|
||||
ImageGenerationPlan::FalFlux1Schnell(plan) => {
|
||||
estimate_image_cost_fal_flux_1_schnell(plan)
|
||||
}
|
||||
ImageGenerationPlan::FalNanoBanana(plan) => {
|
||||
estimate_image_cost_fal_nano_banana(plan)
|
||||
}
|
||||
ImageGenerationPlan::FalNanaBanana2(plan) => {
|
||||
estimate_image_cost_fal_nano_banana_2(plan)
|
||||
}
|
||||
ImageGenerationPlan::FalNanaBananaPro(plan) => {
|
||||
estimate_image_cost_fal_nano_banana_pro(plan)
|
||||
}
|
||||
|
||||
+4
@@ -1 +1,5 @@
|
||||
pub mod plan_generate_image_fal_flux_1_dev;
|
||||
pub mod plan_generate_image_fal_flux_1_schnell;
|
||||
pub mod plan_generate_image_fal_nano_banana;
|
||||
pub mod plan_generate_image_fal_nano_banana_2;
|
||||
pub mod plan_generate_image_fal_nano_banana_pro;
|
||||
|
||||
+158
@@ -0,0 +1,158 @@
|
||||
use crate::api::common_aspect_ratio::CommonAspectRatio;
|
||||
use crate::api::image_list_ref::ImageListRef;
|
||||
use crate::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
|
||||
use crate::errors::artcraft_router_error::ArtcraftRouterError;
|
||||
use crate::errors::client_error::ClientError;
|
||||
use crate::generate::generate_image::generate_image_request::GenerateImageRequest;
|
||||
use crate::generate::generate_image::image_generation_plan::ImageGenerationPlan;
|
||||
use fal_client::requests::webhook::image::edit::enqueue_flux_1_dev_edit_image_webhook::Flux1DevEditImageNumImages;
|
||||
use fal_client::requests::webhook::image::text::enqueue_flux_1_dev_text_to_image_webhook::{
|
||||
Flux1DevAspectRatio, Flux1DevNumImages,
|
||||
};
|
||||
|
||||
/// Shared num-images type for both t2i and edit modes.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum FalFlux1DevNumImages {
|
||||
One,
|
||||
Two,
|
||||
Three,
|
||||
Four,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PlanFalFlux1Dev<'a> {
|
||||
pub prompt: Option<&'a str>,
|
||||
/// Image URL for editing. None = text-to-image mode.
|
||||
pub maybe_image_url: Option<String>,
|
||||
pub aspect_ratio: Flux1DevAspectRatio,
|
||||
pub num_images: FalFlux1DevNumImages,
|
||||
}
|
||||
|
||||
pub fn plan_generate_image_fal_flux_1_dev<'a>(
|
||||
request: &'a GenerateImageRequest<'a>,
|
||||
) -> Result<ImageGenerationPlan<'a>, ArtcraftRouterError> {
|
||||
let strategy = request.request_mismatch_mitigation_strategy;
|
||||
let maybe_image_url = resolve_single_image_url(request.image_inputs)?;
|
||||
let aspect_ratio = plan_aspect_ratio(request.aspect_ratio, strategy)?;
|
||||
let num_images = plan_num_images(request.image_batch_count, strategy)?;
|
||||
|
||||
Ok(ImageGenerationPlan::FalFlux1Dev(PlanFalFlux1Dev {
|
||||
prompt: request.prompt,
|
||||
maybe_image_url,
|
||||
aspect_ratio,
|
||||
num_images,
|
||||
}))
|
||||
}
|
||||
|
||||
fn resolve_single_image_url(
|
||||
image_inputs: Option<ImageListRef<'_>>,
|
||||
) -> Result<Option<String>, ArtcraftRouterError> {
|
||||
match image_inputs {
|
||||
None => Ok(None),
|
||||
Some(ImageListRef::Urls(urls)) => {
|
||||
if urls.is_empty() {
|
||||
Ok(None)
|
||||
} else if urls.len() == 1 {
|
||||
Ok(Some(urls[0].clone()))
|
||||
} else {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "image_inputs",
|
||||
value: format!("Flux 1 Dev image-to-image supports exactly 1 image, got {}", urls.len()),
|
||||
}))
|
||||
}
|
||||
}
|
||||
Some(ImageListRef::MediaFileTokens(_)) => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::FalOnlySupportsUrls))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn plan_aspect_ratio(
|
||||
aspect_ratio: Option<CommonAspectRatio>,
|
||||
strategy: RequestMismatchMitigationStrategy,
|
||||
) -> Result<Flux1DevAspectRatio, ArtcraftRouterError> {
|
||||
use Flux1DevAspectRatio as Ar;
|
||||
match aspect_ratio {
|
||||
None => Ok(Ar::Square),
|
||||
|
||||
Some(CommonAspectRatio::Auto)
|
||||
| Some(CommonAspectRatio::Auto2k)
|
||||
| Some(CommonAspectRatio::Auto4k) => Ok(Ar::Square),
|
||||
|
||||
Some(CommonAspectRatio::Square) => Ok(Ar::Square),
|
||||
Some(CommonAspectRatio::SquareHd) => Ok(Ar::SquareHd),
|
||||
Some(CommonAspectRatio::WideFourByThree) => Ok(Ar::LandscapeFourByThree),
|
||||
Some(CommonAspectRatio::WideSixteenByNine) | Some(CommonAspectRatio::Wide) => Ok(Ar::LandscapeSixteenByNine),
|
||||
Some(CommonAspectRatio::TallThreeByFour) => Ok(Ar::PortraitThreeByFour),
|
||||
Some(CommonAspectRatio::TallNineBySixteen) | Some(CommonAspectRatio::Tall) => Ok(Ar::PortraitNineBySixteen),
|
||||
|
||||
Some(CommonAspectRatio::WideFiveByFour) => Ok(Ar::LandscapeFourByThree),
|
||||
Some(CommonAspectRatio::WideThreeByTwo) => Ok(Ar::LandscapeFourByThree),
|
||||
Some(CommonAspectRatio::WideTwentyOneByNine) => Ok(Ar::LandscapeSixteenByNine),
|
||||
Some(CommonAspectRatio::TallFourByFive) => Ok(Ar::PortraitThreeByFour),
|
||||
Some(CommonAspectRatio::TallTwoByThree) => Ok(Ar::PortraitThreeByFour),
|
||||
Some(CommonAspectRatio::TallNineByTwentyOne) => Ok(Ar::PortraitNineBySixteen),
|
||||
|
||||
Some(unsupported) => match strategy {
|
||||
RequestMismatchMitigationStrategy::ErrorOut => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "aspect_ratio",
|
||||
value: format!("{:?}", unsupported),
|
||||
}))
|
||||
}
|
||||
_ => Ok(Ar::Square),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn plan_num_images(
|
||||
image_batch_count: Option<u16>,
|
||||
strategy: RequestMismatchMitigationStrategy,
|
||||
) -> Result<FalFlux1DevNumImages, ArtcraftRouterError> {
|
||||
let count = image_batch_count.unwrap_or(1);
|
||||
match count {
|
||||
0 => Err(ArtcraftRouterError::Client(ClientError::UserRequestedZeroGenerations)),
|
||||
1 => Ok(FalFlux1DevNumImages::One),
|
||||
2 => Ok(FalFlux1DevNumImages::Two),
|
||||
3 => Ok(FalFlux1DevNumImages::Three),
|
||||
4 => Ok(FalFlux1DevNumImages::Four),
|
||||
_ => match strategy {
|
||||
RequestMismatchMitigationStrategy::ErrorOut => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "image_batch_count",
|
||||
value: format!("{}", count),
|
||||
}))
|
||||
}
|
||||
_ => Ok(FalFlux1DevNumImages::Four),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
impl FalFlux1DevNumImages {
|
||||
pub fn to_t2i(self) -> Flux1DevNumImages {
|
||||
match self {
|
||||
Self::One => Flux1DevNumImages::One,
|
||||
Self::Two => Flux1DevNumImages::Two,
|
||||
Self::Three => Flux1DevNumImages::Three,
|
||||
Self::Four => Flux1DevNumImages::Four,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_edit(self) -> Flux1DevEditImageNumImages {
|
||||
match self {
|
||||
Self::One => Flux1DevEditImageNumImages::One,
|
||||
Self::Two => Flux1DevEditImageNumImages::Two,
|
||||
Self::Three => Flux1DevEditImageNumImages::Three,
|
||||
Self::Four => Flux1DevEditImageNumImages::Four,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_u64(self) -> u64 {
|
||||
match self {
|
||||
Self::One => 1,
|
||||
Self::Two => 2,
|
||||
Self::Three => 3,
|
||||
Self::Four => 4,
|
||||
}
|
||||
}
|
||||
}
|
||||
+169
@@ -0,0 +1,169 @@
|
||||
use crate::api::common_aspect_ratio::CommonAspectRatio;
|
||||
use crate::api::image_list_ref::ImageListRef;
|
||||
use crate::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
|
||||
use crate::errors::artcraft_router_error::ArtcraftRouterError;
|
||||
use crate::errors::client_error::ClientError;
|
||||
use crate::generate::generate_image::generate_image_request::GenerateImageRequest;
|
||||
use crate::generate::generate_image::image_generation_plan::ImageGenerationPlan;
|
||||
use fal_client::requests::webhook::image::edit::enqueue_flux_1_schnell_edit_image_webhook::{
|
||||
Flux1SchnellEditImageNumImages, Flux1SchnellEditImageSize,
|
||||
};
|
||||
use fal_client::requests::webhook::image::text::enqueue_flux_1_schnell_text_to_image_webhook::{
|
||||
Flux1SchnellAspectRatio, Flux1SchnellNumImages,
|
||||
};
|
||||
|
||||
/// Shared num-images type for both t2i and edit modes.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum FalFlux1SchnellNumImages {
|
||||
One,
|
||||
Two,
|
||||
Three,
|
||||
Four,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PlanFalFlux1Schnell<'a> {
|
||||
pub prompt: Option<&'a str>,
|
||||
/// Image URL for editing. None = text-to-image mode.
|
||||
pub maybe_image_url: Option<String>,
|
||||
pub t2i_aspect_ratio: Flux1SchnellAspectRatio,
|
||||
pub edit_image_size: Option<Flux1SchnellEditImageSize>,
|
||||
pub num_images: FalFlux1SchnellNumImages,
|
||||
}
|
||||
|
||||
pub fn plan_generate_image_fal_flux_1_schnell<'a>(
|
||||
request: &'a GenerateImageRequest<'a>,
|
||||
) -> Result<ImageGenerationPlan<'a>, ArtcraftRouterError> {
|
||||
let strategy = request.request_mismatch_mitigation_strategy;
|
||||
let maybe_image_url = resolve_single_image_url(request.image_inputs)?;
|
||||
let t2i_aspect_ratio = plan_t2i_aspect_ratio(request.aspect_ratio, strategy)?;
|
||||
let edit_image_size = plan_edit_image_size(request.aspect_ratio);
|
||||
let num_images = plan_num_images(request.image_batch_count, strategy)?;
|
||||
|
||||
Ok(ImageGenerationPlan::FalFlux1Schnell(PlanFalFlux1Schnell {
|
||||
prompt: request.prompt,
|
||||
maybe_image_url,
|
||||
t2i_aspect_ratio,
|
||||
edit_image_size,
|
||||
num_images,
|
||||
}))
|
||||
}
|
||||
|
||||
fn resolve_single_image_url(
|
||||
image_inputs: Option<ImageListRef<'_>>,
|
||||
) -> Result<Option<String>, ArtcraftRouterError> {
|
||||
match image_inputs {
|
||||
None => Ok(None),
|
||||
Some(ImageListRef::Urls(urls)) => {
|
||||
if urls.is_empty() {
|
||||
Ok(None)
|
||||
} else if urls.len() == 1 {
|
||||
Ok(Some(urls[0].clone()))
|
||||
} else {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "image_inputs",
|
||||
value: format!("Flux 1 Schnell redux supports exactly 1 image, got {}", urls.len()),
|
||||
}))
|
||||
}
|
||||
}
|
||||
Some(ImageListRef::MediaFileTokens(_)) => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::FalOnlySupportsUrls))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn plan_t2i_aspect_ratio(
|
||||
aspect_ratio: Option<CommonAspectRatio>,
|
||||
strategy: RequestMismatchMitigationStrategy,
|
||||
) -> Result<Flux1SchnellAspectRatio, ArtcraftRouterError> {
|
||||
use Flux1SchnellAspectRatio as Ar;
|
||||
match aspect_ratio {
|
||||
None => Ok(Ar::Square),
|
||||
Some(CommonAspectRatio::Auto) | Some(CommonAspectRatio::Auto2k) | Some(CommonAspectRatio::Auto4k) => Ok(Ar::Square),
|
||||
Some(CommonAspectRatio::Square) => Ok(Ar::Square),
|
||||
Some(CommonAspectRatio::SquareHd) => Ok(Ar::SquareHd),
|
||||
Some(CommonAspectRatio::WideFourByThree) => Ok(Ar::LandscapeFourByThree),
|
||||
Some(CommonAspectRatio::WideSixteenByNine) | Some(CommonAspectRatio::Wide) => Ok(Ar::LandscapeSixteenByNine),
|
||||
Some(CommonAspectRatio::TallThreeByFour) => Ok(Ar::PortraitThreeByFour),
|
||||
Some(CommonAspectRatio::TallNineBySixteen) | Some(CommonAspectRatio::Tall) => Ok(Ar::PortraitNineBySixteen),
|
||||
Some(CommonAspectRatio::WideFiveByFour) | Some(CommonAspectRatio::WideThreeByTwo) => Ok(Ar::LandscapeFourByThree),
|
||||
Some(CommonAspectRatio::WideTwentyOneByNine) => Ok(Ar::LandscapeSixteenByNine),
|
||||
Some(CommonAspectRatio::TallFourByFive) | Some(CommonAspectRatio::TallTwoByThree) => Ok(Ar::PortraitThreeByFour),
|
||||
Some(CommonAspectRatio::TallNineByTwentyOne) => Ok(Ar::PortraitNineBySixteen),
|
||||
Some(unsupported) => match strategy {
|
||||
RequestMismatchMitigationStrategy::ErrorOut => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "aspect_ratio",
|
||||
value: format!("{:?}", unsupported),
|
||||
}))
|
||||
}
|
||||
_ => Ok(Ar::Square),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn plan_edit_image_size(aspect_ratio: Option<CommonAspectRatio>) -> Option<Flux1SchnellEditImageSize> {
|
||||
use Flux1SchnellEditImageSize as Size;
|
||||
match aspect_ratio {
|
||||
None => None,
|
||||
Some(CommonAspectRatio::Square) => Some(Size::Square),
|
||||
Some(CommonAspectRatio::SquareHd) => Some(Size::SquareHd),
|
||||
Some(CommonAspectRatio::WideFourByThree) => Some(Size::LandscapeFourByThree),
|
||||
Some(CommonAspectRatio::WideSixteenByNine) | Some(CommonAspectRatio::Wide) => Some(Size::LandscapeSixteenByNine),
|
||||
Some(CommonAspectRatio::TallThreeByFour) => Some(Size::PortraitThreeByFour),
|
||||
Some(CommonAspectRatio::TallNineBySixteen) | Some(CommonAspectRatio::Tall) => Some(Size::PortraitNineBySixteen),
|
||||
_ => Some(Size::Square),
|
||||
}
|
||||
}
|
||||
|
||||
fn plan_num_images(
|
||||
image_batch_count: Option<u16>,
|
||||
strategy: RequestMismatchMitigationStrategy,
|
||||
) -> Result<FalFlux1SchnellNumImages, ArtcraftRouterError> {
|
||||
let count = image_batch_count.unwrap_or(1);
|
||||
match count {
|
||||
0 => Err(ArtcraftRouterError::Client(ClientError::UserRequestedZeroGenerations)),
|
||||
1 => Ok(FalFlux1SchnellNumImages::One),
|
||||
2 => Ok(FalFlux1SchnellNumImages::Two),
|
||||
3 => Ok(FalFlux1SchnellNumImages::Three),
|
||||
4 => Ok(FalFlux1SchnellNumImages::Four),
|
||||
_ => match strategy {
|
||||
RequestMismatchMitigationStrategy::ErrorOut => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "image_batch_count",
|
||||
value: format!("{}", count),
|
||||
}))
|
||||
}
|
||||
_ => Ok(FalFlux1SchnellNumImages::Four),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
impl FalFlux1SchnellNumImages {
|
||||
pub fn to_t2i(self) -> Flux1SchnellNumImages {
|
||||
match self {
|
||||
Self::One => Flux1SchnellNumImages::One,
|
||||
Self::Two => Flux1SchnellNumImages::Two,
|
||||
Self::Three => Flux1SchnellNumImages::Three,
|
||||
Self::Four => Flux1SchnellNumImages::Four,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_edit(self) -> Flux1SchnellEditImageNumImages {
|
||||
match self {
|
||||
Self::One => Flux1SchnellEditImageNumImages::One,
|
||||
Self::Two => Flux1SchnellEditImageNumImages::Two,
|
||||
Self::Three => Flux1SchnellEditImageNumImages::Three,
|
||||
Self::Four => Flux1SchnellEditImageNumImages::Four,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_u64(self) -> u64 {
|
||||
match self {
|
||||
Self::One => 1,
|
||||
Self::Two => 2,
|
||||
Self::Three => 3,
|
||||
Self::Four => 4,
|
||||
}
|
||||
}
|
||||
}
|
||||
+174
@@ -0,0 +1,174 @@
|
||||
use crate::api::common_aspect_ratio::CommonAspectRatio;
|
||||
use crate::api::image_list_ref::ImageListRef;
|
||||
use crate::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
|
||||
use crate::errors::artcraft_router_error::ArtcraftRouterError;
|
||||
use crate::errors::client_error::ClientError;
|
||||
use crate::generate::generate_image::generate_image_request::GenerateImageRequest;
|
||||
use crate::generate::generate_image::image_generation_plan::ImageGenerationPlan;
|
||||
use fal_client::requests::webhook::image::edit::enqueue_gemini_25_flash_edit_webhook::Gemini25FlashEditAspectRatio;
|
||||
use fal_client::requests::webhook::image::text::enqueue_gemini_25_flash_text_to_image_webhook::Gemini25FlashTextToImageAspectRatio;
|
||||
|
||||
/// Intermediate num-images type shared between t2i and edit enums.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum FalNbNumImages {
|
||||
One,
|
||||
Two,
|
||||
Three,
|
||||
Four,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PlanFalNanoBanana<'a> {
|
||||
pub prompt: Option<&'a str>,
|
||||
/// Image URLs for editing. Empty vec = text-to-image mode.
|
||||
pub image_urls: Vec<String>,
|
||||
/// Pre-resolved aspect ratio for text-to-image mode.
|
||||
pub t2i_aspect_ratio: Option<Gemini25FlashTextToImageAspectRatio>,
|
||||
/// Pre-resolved aspect ratio for image-edit mode.
|
||||
pub edit_aspect_ratio: Option<Gemini25FlashEditAspectRatio>,
|
||||
/// Pre-resolved number of images (1–4, shared for both modes).
|
||||
pub num_images: FalNbNumImages,
|
||||
}
|
||||
|
||||
pub fn plan_generate_image_fal_nano_banana<'a>(
|
||||
request: &'a GenerateImageRequest<'a>,
|
||||
) -> Result<ImageGenerationPlan<'a>, ArtcraftRouterError> {
|
||||
let strategy = request.request_mismatch_mitigation_strategy;
|
||||
let is_edit_mode = request.image_inputs.is_some();
|
||||
let image_urls = resolve_image_list_ref(request.image_inputs)?;
|
||||
let t2i_aspect_ratio = plan_t2i_aspect_ratio(request.aspect_ratio, strategy)?;
|
||||
let edit_aspect_ratio = plan_edit_aspect_ratio(request.aspect_ratio, is_edit_mode, strategy)?;
|
||||
let num_images = plan_num_images(request.image_batch_count, strategy)?;
|
||||
|
||||
Ok(ImageGenerationPlan::FalNanoBanana(PlanFalNanoBanana {
|
||||
prompt: request.prompt,
|
||||
image_urls,
|
||||
t2i_aspect_ratio,
|
||||
edit_aspect_ratio,
|
||||
num_images,
|
||||
}))
|
||||
}
|
||||
|
||||
fn resolve_image_list_ref(
|
||||
image_list_ref: Option<ImageListRef<'_>>,
|
||||
) -> Result<Vec<String>, ArtcraftRouterError> {
|
||||
match image_list_ref {
|
||||
None => Ok(vec![]),
|
||||
Some(ImageListRef::Urls(urls)) => Ok(urls.clone()),
|
||||
Some(ImageListRef::MediaFileTokens(_)) => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::FalOnlySupportsUrls))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Gemini 25 Flash text-to-image supported aspect ratios:
|
||||
// 1:1, 5:4, 4:3, 3:2, 16:9, 21:9, 4:5, 3:4, 2:3, 9:16
|
||||
// (no Auto for text-to-image)
|
||||
fn plan_t2i_aspect_ratio(
|
||||
aspect_ratio: Option<CommonAspectRatio>,
|
||||
strategy: RequestMismatchMitigationStrategy,
|
||||
) -> Result<Option<Gemini25FlashTextToImageAspectRatio>, ArtcraftRouterError> {
|
||||
use Gemini25FlashTextToImageAspectRatio as T2iAr;
|
||||
match aspect_ratio {
|
||||
None => Ok(None),
|
||||
|
||||
// Auto not valid for text-to-image; fall back to square
|
||||
Some(CommonAspectRatio::Auto)
|
||||
| Some(CommonAspectRatio::Auto2k)
|
||||
| Some(CommonAspectRatio::Auto4k) => Ok(Some(T2iAr::OneByOne)),
|
||||
|
||||
Some(CommonAspectRatio::Square) | Some(CommonAspectRatio::SquareHd) => Ok(Some(T2iAr::OneByOne)),
|
||||
Some(CommonAspectRatio::WideFiveByFour) => Ok(Some(T2iAr::FiveByFour)),
|
||||
Some(CommonAspectRatio::WideFourByThree) => Ok(Some(T2iAr::FourByThree)),
|
||||
Some(CommonAspectRatio::WideThreeByTwo) => Ok(Some(T2iAr::ThreeByTwo)),
|
||||
Some(CommonAspectRatio::WideSixteenByNine) | Some(CommonAspectRatio::Wide) => Ok(Some(T2iAr::SixteenByNine)),
|
||||
Some(CommonAspectRatio::WideTwentyOneByNine) => Ok(Some(T2iAr::TwentyOneByNine)),
|
||||
Some(CommonAspectRatio::TallFourByFive) => Ok(Some(T2iAr::FourByFive)),
|
||||
Some(CommonAspectRatio::TallThreeByFour) => Ok(Some(T2iAr::ThreeByFour)),
|
||||
Some(CommonAspectRatio::TallTwoByThree) => Ok(Some(T2iAr::TwoByThree)),
|
||||
Some(CommonAspectRatio::TallNineBySixteen) | Some(CommonAspectRatio::Tall) => Ok(Some(T2iAr::NineBySixteen)),
|
||||
|
||||
Some(unsupported) => match strategy {
|
||||
RequestMismatchMitigationStrategy::ErrorOut => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "aspect_ratio",
|
||||
value: format!("{:?}", unsupported),
|
||||
}))
|
||||
}
|
||||
RequestMismatchMitigationStrategy::PayMoreUpgrade
|
||||
| RequestMismatchMitigationStrategy::PayLessDowngrade => {
|
||||
Ok(Some(T2iAr::NineBySixteen))
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Gemini 25 Flash image-edit supported aspect ratios: same as t2i plus Auto
|
||||
fn plan_edit_aspect_ratio(
|
||||
aspect_ratio: Option<CommonAspectRatio>,
|
||||
is_edit_mode: bool,
|
||||
strategy: RequestMismatchMitigationStrategy,
|
||||
) -> Result<Option<Gemini25FlashEditAspectRatio>, ArtcraftRouterError> {
|
||||
use Gemini25FlashEditAspectRatio as EditAr;
|
||||
match aspect_ratio {
|
||||
None => Ok(None),
|
||||
|
||||
Some(CommonAspectRatio::Auto)
|
||||
| Some(CommonAspectRatio::Auto2k)
|
||||
| Some(CommonAspectRatio::Auto4k) => {
|
||||
if is_edit_mode {
|
||||
Ok(Some(EditAr::Auto))
|
||||
} else {
|
||||
Ok(Some(EditAr::OneByOne))
|
||||
}
|
||||
}
|
||||
|
||||
Some(CommonAspectRatio::Square) | Some(CommonAspectRatio::SquareHd) => Ok(Some(EditAr::OneByOne)),
|
||||
Some(CommonAspectRatio::WideFiveByFour) => Ok(Some(EditAr::FiveByFour)),
|
||||
Some(CommonAspectRatio::WideFourByThree) => Ok(Some(EditAr::FourByThree)),
|
||||
Some(CommonAspectRatio::WideThreeByTwo) => Ok(Some(EditAr::ThreeByTwo)),
|
||||
Some(CommonAspectRatio::WideSixteenByNine) | Some(CommonAspectRatio::Wide) => Ok(Some(EditAr::SixteenByNine)),
|
||||
Some(CommonAspectRatio::WideTwentyOneByNine) => Ok(Some(EditAr::TwentyOneByNine)),
|
||||
Some(CommonAspectRatio::TallFourByFive) => Ok(Some(EditAr::FourByFive)),
|
||||
Some(CommonAspectRatio::TallThreeByFour) => Ok(Some(EditAr::ThreeByFour)),
|
||||
Some(CommonAspectRatio::TallTwoByThree) => Ok(Some(EditAr::TwoByThree)),
|
||||
Some(CommonAspectRatio::TallNineBySixteen) | Some(CommonAspectRatio::Tall) => Ok(Some(EditAr::NineBySixteen)),
|
||||
|
||||
Some(unsupported) => match strategy {
|
||||
RequestMismatchMitigationStrategy::ErrorOut => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "aspect_ratio",
|
||||
value: format!("{:?}", unsupported),
|
||||
}))
|
||||
}
|
||||
RequestMismatchMitigationStrategy::PayMoreUpgrade
|
||||
| RequestMismatchMitigationStrategy::PayLessDowngrade => {
|
||||
Ok(Some(EditAr::NineBySixteen))
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn plan_num_images(
|
||||
image_batch_count: Option<u16>,
|
||||
strategy: RequestMismatchMitigationStrategy,
|
||||
) -> Result<FalNbNumImages, ArtcraftRouterError> {
|
||||
let count = image_batch_count.unwrap_or(1);
|
||||
match count {
|
||||
0 => Err(ArtcraftRouterError::Client(ClientError::UserRequestedZeroGenerations)),
|
||||
1 => Ok(FalNbNumImages::One),
|
||||
2 => Ok(FalNbNumImages::Two),
|
||||
3 => Ok(FalNbNumImages::Three),
|
||||
4 => Ok(FalNbNumImages::Four),
|
||||
_ => match strategy {
|
||||
RequestMismatchMitigationStrategy::ErrorOut => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "image_batch_count",
|
||||
value: format!("{}", count),
|
||||
}))
|
||||
}
|
||||
RequestMismatchMitigationStrategy::PayMoreUpgrade
|
||||
| RequestMismatchMitigationStrategy::PayLessDowngrade => Ok(FalNbNumImages::Four),
|
||||
},
|
||||
}
|
||||
}
|
||||
+202
@@ -0,0 +1,202 @@
|
||||
use crate::api::common_aspect_ratio::CommonAspectRatio;
|
||||
use crate::api::common_resolution::CommonResolution;
|
||||
use crate::api::image_list_ref::ImageListRef;
|
||||
use crate::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
|
||||
use crate::errors::artcraft_router_error::ArtcraftRouterError;
|
||||
use crate::errors::client_error::ClientError;
|
||||
use crate::generate::generate_image::generate_image_request::GenerateImageRequest;
|
||||
use crate::generate::generate_image::image_generation_plan::ImageGenerationPlan;
|
||||
use fal_client::requests::webhook::image::edit::enqueue_nano_banana_2_edit_image_webhook::EnqueueNanoBanana2EditImageAspectRatio;
|
||||
use fal_client::requests::webhook::image::text::enqueue_nano_banana_2_text_to_image_webhook::EnqueueNanoBanana2TextToImageAspectRatio;
|
||||
|
||||
/// Intermediate resolution type shared between t2i and edit enums.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum FalNb2Resolution {
|
||||
HalfK,
|
||||
OneK,
|
||||
TwoK,
|
||||
FourK,
|
||||
}
|
||||
|
||||
/// Intermediate num-images type shared between t2i and edit enums.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum FalNb2NumImages {
|
||||
One,
|
||||
Two,
|
||||
Three,
|
||||
Four,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PlanFalNanaBanana2<'a> {
|
||||
pub prompt: Option<&'a str>,
|
||||
/// Image URLs for editing. Empty vec = text-to-image mode.
|
||||
pub image_urls: Vec<String>,
|
||||
/// Pre-resolved aspect ratio for text-to-image mode.
|
||||
pub t2i_aspect_ratio: Option<EnqueueNanoBanana2TextToImageAspectRatio>,
|
||||
/// Pre-resolved aspect ratio for image-edit mode.
|
||||
pub edit_aspect_ratio: Option<EnqueueNanoBanana2EditImageAspectRatio>,
|
||||
/// Pre-resolved resolution (shared for both modes).
|
||||
pub resolution: Option<FalNb2Resolution>,
|
||||
/// Pre-resolved number of images (1–4, shared for both modes).
|
||||
pub num_images: FalNb2NumImages,
|
||||
}
|
||||
|
||||
pub fn plan_generate_image_fal_nano_banana_2<'a>(
|
||||
request: &'a GenerateImageRequest<'a>,
|
||||
) -> Result<ImageGenerationPlan<'a>, ArtcraftRouterError> {
|
||||
let strategy = request.request_mismatch_mitigation_strategy;
|
||||
let is_edit_mode = request.image_inputs.is_some();
|
||||
let image_urls = resolve_image_list_ref(request.image_inputs)?;
|
||||
let t2i_aspect_ratio = plan_t2i_aspect_ratio(request.aspect_ratio, strategy)?;
|
||||
let edit_aspect_ratio = plan_edit_aspect_ratio(request.aspect_ratio, is_edit_mode, strategy)?;
|
||||
let resolution = plan_resolution(request.resolution)?;
|
||||
let num_images = plan_num_images(request.image_batch_count, strategy)?;
|
||||
|
||||
Ok(ImageGenerationPlan::FalNanaBanana2(PlanFalNanaBanana2 {
|
||||
prompt: request.prompt,
|
||||
image_urls,
|
||||
t2i_aspect_ratio,
|
||||
edit_aspect_ratio,
|
||||
resolution,
|
||||
num_images,
|
||||
}))
|
||||
}
|
||||
|
||||
fn resolve_image_list_ref(
|
||||
image_list_ref: Option<ImageListRef<'_>>,
|
||||
) -> Result<Vec<String>, ArtcraftRouterError> {
|
||||
match image_list_ref {
|
||||
None => Ok(vec![]),
|
||||
Some(ImageListRef::Urls(urls)) => Ok(urls.clone()),
|
||||
Some(ImageListRef::MediaFileTokens(_)) => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::FalOnlySupportsUrls))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// NB2 text-to-image supported aspect ratios:
|
||||
// auto, 1:1, 5:4, 4:3, 3:2, 16:9, 21:9, 4:5, 3:4, 2:3, 9:16
|
||||
fn plan_t2i_aspect_ratio(
|
||||
aspect_ratio: Option<CommonAspectRatio>,
|
||||
strategy: RequestMismatchMitigationStrategy,
|
||||
) -> Result<Option<EnqueueNanoBanana2TextToImageAspectRatio>, ArtcraftRouterError> {
|
||||
use EnqueueNanoBanana2TextToImageAspectRatio as T2iAr;
|
||||
match aspect_ratio {
|
||||
None => Ok(None),
|
||||
|
||||
Some(CommonAspectRatio::Auto)
|
||||
| Some(CommonAspectRatio::Auto2k)
|
||||
| Some(CommonAspectRatio::Auto4k) => Ok(Some(T2iAr::Auto)),
|
||||
|
||||
Some(CommonAspectRatio::Square) | Some(CommonAspectRatio::SquareHd) => Ok(Some(T2iAr::OneByOne)),
|
||||
Some(CommonAspectRatio::WideFiveByFour) => Ok(Some(T2iAr::FiveByFour)),
|
||||
Some(CommonAspectRatio::WideFourByThree) => Ok(Some(T2iAr::FourByThree)),
|
||||
Some(CommonAspectRatio::WideThreeByTwo) => Ok(Some(T2iAr::ThreeByTwo)),
|
||||
Some(CommonAspectRatio::WideSixteenByNine) | Some(CommonAspectRatio::Wide) => Ok(Some(T2iAr::SixteenByNine)),
|
||||
Some(CommonAspectRatio::WideTwentyOneByNine) => Ok(Some(T2iAr::TwentyOneByNine)),
|
||||
Some(CommonAspectRatio::TallFourByFive) => Ok(Some(T2iAr::FourByFive)),
|
||||
Some(CommonAspectRatio::TallThreeByFour) => Ok(Some(T2iAr::ThreeByFour)),
|
||||
Some(CommonAspectRatio::TallTwoByThree) => Ok(Some(T2iAr::TwoByThree)),
|
||||
Some(CommonAspectRatio::TallNineBySixteen) | Some(CommonAspectRatio::Tall) => Ok(Some(T2iAr::NineBySixteen)),
|
||||
|
||||
Some(unsupported) => match strategy {
|
||||
RequestMismatchMitigationStrategy::ErrorOut => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "aspect_ratio",
|
||||
value: format!("{:?}", unsupported),
|
||||
}))
|
||||
}
|
||||
RequestMismatchMitigationStrategy::PayMoreUpgrade
|
||||
| RequestMismatchMitigationStrategy::PayLessDowngrade => {
|
||||
Ok(Some(T2iAr::NineBySixteen))
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// NB2 image-edit supported aspect ratios: same as t2i plus Auto
|
||||
fn plan_edit_aspect_ratio(
|
||||
aspect_ratio: Option<CommonAspectRatio>,
|
||||
is_edit_mode: bool,
|
||||
strategy: RequestMismatchMitigationStrategy,
|
||||
) -> Result<Option<EnqueueNanoBanana2EditImageAspectRatio>, ArtcraftRouterError> {
|
||||
use EnqueueNanoBanana2EditImageAspectRatio as EditAr;
|
||||
match aspect_ratio {
|
||||
None => Ok(None),
|
||||
|
||||
Some(CommonAspectRatio::Auto)
|
||||
| Some(CommonAspectRatio::Auto2k)
|
||||
| Some(CommonAspectRatio::Auto4k) => {
|
||||
if is_edit_mode {
|
||||
Ok(Some(EditAr::Auto))
|
||||
} else {
|
||||
Ok(Some(EditAr::OneByOne))
|
||||
}
|
||||
}
|
||||
|
||||
Some(CommonAspectRatio::Square) | Some(CommonAspectRatio::SquareHd) => Ok(Some(EditAr::OneByOne)),
|
||||
Some(CommonAspectRatio::WideFiveByFour) => Ok(Some(EditAr::FiveByFour)),
|
||||
Some(CommonAspectRatio::WideFourByThree) => Ok(Some(EditAr::FourByThree)),
|
||||
Some(CommonAspectRatio::WideThreeByTwo) => Ok(Some(EditAr::ThreeByTwo)),
|
||||
Some(CommonAspectRatio::WideSixteenByNine) | Some(CommonAspectRatio::Wide) => Ok(Some(EditAr::SixteenByNine)),
|
||||
Some(CommonAspectRatio::WideTwentyOneByNine) => Ok(Some(EditAr::TwentyOneByNine)),
|
||||
Some(CommonAspectRatio::TallFourByFive) => Ok(Some(EditAr::FourByFive)),
|
||||
Some(CommonAspectRatio::TallThreeByFour) => Ok(Some(EditAr::ThreeByFour)),
|
||||
Some(CommonAspectRatio::TallTwoByThree) => Ok(Some(EditAr::TwoByThree)),
|
||||
Some(CommonAspectRatio::TallNineBySixteen) | Some(CommonAspectRatio::Tall) => Ok(Some(EditAr::NineBySixteen)),
|
||||
|
||||
Some(unsupported) => match strategy {
|
||||
RequestMismatchMitigationStrategy::ErrorOut => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "aspect_ratio",
|
||||
value: format!("{:?}", unsupported),
|
||||
}))
|
||||
}
|
||||
RequestMismatchMitigationStrategy::PayMoreUpgrade
|
||||
| RequestMismatchMitigationStrategy::PayLessDowngrade => {
|
||||
Ok(Some(EditAr::NineBySixteen))
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn plan_resolution(
|
||||
resolution: Option<CommonResolution>,
|
||||
) -> Result<Option<FalNb2Resolution>, ArtcraftRouterError> {
|
||||
match resolution {
|
||||
None => Ok(None),
|
||||
Some(CommonResolution::HalfK) => Ok(Some(FalNb2Resolution::HalfK)),
|
||||
Some(CommonResolution::OneK) => Ok(Some(FalNb2Resolution::OneK)),
|
||||
Some(CommonResolution::TwoK) => Ok(Some(FalNb2Resolution::TwoK)),
|
||||
Some(CommonResolution::FourK) => Ok(Some(FalNb2Resolution::FourK)),
|
||||
Some(CommonResolution::ThreeK) => Ok(Some(FalNb2Resolution::TwoK)),
|
||||
Some(CommonResolution::FourEightyP) => Ok(Some(FalNb2Resolution::HalfK)),
|
||||
Some(CommonResolution::SevenTwentyP) => Ok(Some(FalNb2Resolution::OneK)),
|
||||
Some(CommonResolution::TenEightyP) => Ok(Some(FalNb2Resolution::OneK)),
|
||||
}
|
||||
}
|
||||
|
||||
fn plan_num_images(
|
||||
image_batch_count: Option<u16>,
|
||||
strategy: RequestMismatchMitigationStrategy,
|
||||
) -> Result<FalNb2NumImages, ArtcraftRouterError> {
|
||||
let count = image_batch_count.unwrap_or(1);
|
||||
match count {
|
||||
0 => Err(ArtcraftRouterError::Client(ClientError::UserRequestedZeroGenerations)),
|
||||
1 => Ok(FalNb2NumImages::One),
|
||||
2 => Ok(FalNb2NumImages::Two),
|
||||
3 => Ok(FalNb2NumImages::Three),
|
||||
4 => Ok(FalNb2NumImages::Four),
|
||||
_ => match strategy {
|
||||
RequestMismatchMitigationStrategy::ErrorOut => {
|
||||
Err(ArtcraftRouterError::Client(ClientError::ModelDoesNotSupportOption {
|
||||
field: "image_batch_count",
|
||||
value: format!("{}", count),
|
||||
}))
|
||||
}
|
||||
RequestMismatchMitigationStrategy::PayMoreUpgrade
|
||||
| RequestMismatchMitigationStrategy::PayLessDowngrade => Ok(FalNb2NumImages::Four),
|
||||
},
|
||||
}
|
||||
}
|
||||
+1
-1
@@ -250,7 +250,7 @@ mod tests {
|
||||
|
||||
fn base_fal_request() -> GenerateImageRequest<'static> {
|
||||
GenerateImageRequest {
|
||||
model: CommonImageModel::NanaBananaPro,
|
||||
model: CommonImageModel::NanoBananaPro,
|
||||
provider: Provider::Fal,
|
||||
prompt: Some("a cat in space"),
|
||||
image_inputs: None,
|
||||
|
||||
+3
-3
@@ -1,9 +1,8 @@
|
||||
use tokens::tokens::characters::CharacterToken;
|
||||
use crate::api::audio_list_ref::AudioListRef;
|
||||
use crate::api::character_list_ref::CharacterListRef;
|
||||
use crate::api::common_aspect_ratio::CommonAspectRatio;
|
||||
use crate::api::common_resolution::CommonResolution;
|
||||
use crate::api::common_video_model::CommonVideoModel;
|
||||
use crate::api::audio_list_ref::AudioListRef;
|
||||
use crate::api::character_list_ref::CharacterListRef;
|
||||
use crate::api::image_list_ref::ImageListRef;
|
||||
use crate::api::image_ref::ImageRef;
|
||||
use crate::api::provider::Provider;
|
||||
@@ -21,6 +20,7 @@ use crate::generate::generate_video::video_generation_plan::VideoGenerationPlan;
|
||||
/// Plan to either (1) generate a video or (2) determine how much it costs to generate that video.
|
||||
/// This works across multiple providers by shaping a generic "GenerateVideoRequest" into a provider-specific plan.
|
||||
/// That plan can then be used to return a cost estimate for that given provider or return a struct that can be used to send a real generation request.
|
||||
#[derive(Debug)]
|
||||
pub struct GenerateVideoRequest<'a> {
|
||||
/// Which model to use.
|
||||
pub model: CommonVideoModel,
|
||||
|
||||
@@ -23,7 +23,7 @@ pub fn get_artcraft_client() -> RouterClient {
|
||||
|
||||
pub fn base_image_request() -> GenerateImageRequest<'static> {
|
||||
GenerateImageRequest {
|
||||
model: CommonImageModel::NanaBananaPro,
|
||||
model: CommonImageModel::NanoBananaPro,
|
||||
provider: Provider::Artcraft,
|
||||
prompt: Some("a cat in space"),
|
||||
image_inputs: None,
|
||||
@@ -49,7 +49,7 @@ pub fn get_fal_client() -> RouterClient {
|
||||
|
||||
pub fn base_fal_image_request() -> GenerateImageRequest<'static> {
|
||||
GenerateImageRequest {
|
||||
model: CommonImageModel::NanaBananaPro,
|
||||
model: CommonImageModel::NanoBananaPro,
|
||||
provider: Provider::Fal,
|
||||
prompt: Some("a cat in space"),
|
||||
image_inputs: None,
|
||||
@@ -121,7 +121,7 @@ pub fn base_seedream_5_lite_image_request() -> GenerateImageRequest<'static> {
|
||||
|
||||
pub fn base_nano_banana_2_image_request() -> GenerateImageRequest<'static> {
|
||||
GenerateImageRequest {
|
||||
model: CommonImageModel::NanaBanana2,
|
||||
model: CommonImageModel::NanoBanana2,
|
||||
provider: Provider::Artcraft,
|
||||
prompt: Some("a cat in space"),
|
||||
image_inputs: None,
|
||||
@@ -139,7 +139,7 @@ pub fn base_nano_banana_2_image_request() -> GenerateImageRequest<'static> {
|
||||
|
||||
pub fn base_nano_banana_image_request() -> GenerateImageRequest<'static> {
|
||||
GenerateImageRequest {
|
||||
model: CommonImageModel::NanaBanana,
|
||||
model: CommonImageModel::NanoBanana,
|
||||
provider: Provider::Artcraft,
|
||||
prompt: Some("a cat in space"),
|
||||
image_inputs: None,
|
||||
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
use crate::requests::api::fal_request::FalRequest;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Default)]
|
||||
pub struct Flux1DevEditImageInput {
|
||||
pub prompt: String,
|
||||
|
||||
pub image_url: String,
|
||||
|
||||
/// 0.0 - 1.0
|
||||
/// Default: 0.95
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub strength: Option<f64>,
|
||||
|
||||
/// 1 - 4
|
||||
/// Default: 1
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub num_images: Option<i64>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub seed: Option<i64>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub guidance_scale: Option<f64>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub num_inference_steps: Option<i64>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub enable_safety_checker: Option<bool>,
|
||||
|
||||
/// Options: jpeg, png
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub output_format: Option<String>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub sync_mode: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct Flux1DevEditImageFile {
|
||||
pub url: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct Flux1DevEditImageOutput {
|
||||
pub images: Vec<Flux1DevEditImageFile>,
|
||||
}
|
||||
|
||||
pub fn flux_1_dev_edit_image(
|
||||
params: Flux1DevEditImageInput,
|
||||
) -> FalRequest<Flux1DevEditImageInput, Flux1DevEditImageOutput> {
|
||||
FalRequest::new("fal-ai/flux/dev/image-to-image", params)
|
||||
}
|
||||
+48
@@ -0,0 +1,48 @@
|
||||
use crate::requests::api::fal_request::FalRequest;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Default)]
|
||||
pub struct Flux1SchnellEditImageInput {
|
||||
pub image_url: String,
|
||||
|
||||
/// Options: square_hd, square, portrait_4_3, portrait_16_9, landscape_4_3, landscape_16_9
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub image_size: Option<String>,
|
||||
|
||||
/// 1 - 4
|
||||
/// Default: 1
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub num_images: Option<i64>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub seed: Option<i64>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub num_inference_steps: Option<i64>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub enable_safety_checker: Option<bool>,
|
||||
|
||||
/// Options: jpeg, png
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub output_format: Option<String>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub sync_mode: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct Flux1SchnellEditImageFile {
|
||||
pub url: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct Flux1SchnellEditImageOutput {
|
||||
pub images: Vec<Flux1SchnellEditImageFile>,
|
||||
}
|
||||
|
||||
pub fn flux_1_schnell_edit_image(
|
||||
params: Flux1SchnellEditImageInput,
|
||||
) -> FalRequest<Flux1SchnellEditImageInput, Flux1SchnellEditImageOutput> {
|
||||
FalRequest::new("fal-ai/flux/schnell/redux", params)
|
||||
}
|
||||
@@ -1,3 +1,5 @@
|
||||
pub mod http_flux_1_dev_edit_image;
|
||||
pub mod http_flux_1_schnell_edit_image;
|
||||
pub mod http_flux_pro_kontext_max_edit;
|
||||
pub mod http_gemini_25_flash_edit;
|
||||
pub mod http_gpt_image_1_edit_image;
|
||||
|
||||
+100
@@ -0,0 +1,100 @@
|
||||
use crate::creds::fal_api_key::FalApiKey;
|
||||
use crate::error::classify_fal_error::classify_fal_error;
|
||||
use crate::error::fal_error_plus::FalErrorPlus;
|
||||
use crate::requests::http::image::edit::http_flux_1_dev_edit_image::{flux_1_dev_edit_image, Flux1DevEditImageInput};
|
||||
use crate::requests::api::webhook_response::WebhookResponse;
|
||||
use reqwest::IntoUrl;
|
||||
|
||||
pub struct Flux1DevEditImageArgs<'a, U: IntoUrl> {
|
||||
pub prompt: &'a str,
|
||||
pub image_url: String,
|
||||
pub num_images: Flux1DevEditImageNumImages,
|
||||
|
||||
// Fulfillment
|
||||
pub webhook_url: U,
|
||||
pub api_key: &'a FalApiKey,
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub enum Flux1DevEditImageNumImages {
|
||||
One,
|
||||
Two,
|
||||
Three,
|
||||
Four,
|
||||
}
|
||||
|
||||
pub async fn enqueue_flux_1_dev_edit_image_webhook<U: IntoUrl>(
|
||||
args: Flux1DevEditImageArgs<'_, U>
|
||||
) -> Result<WebhookResponse, FalErrorPlus> {
|
||||
let num_images = match args.num_images {
|
||||
Flux1DevEditImageNumImages::One => 1,
|
||||
Flux1DevEditImageNumImages::Two => 2,
|
||||
Flux1DevEditImageNumImages::Three => 3,
|
||||
Flux1DevEditImageNumImages::Four => 4,
|
||||
};
|
||||
|
||||
let request = Flux1DevEditImageInput {
|
||||
prompt: args.prompt.to_string(),
|
||||
image_url: args.image_url,
|
||||
num_images: Some(num_images),
|
||||
enable_safety_checker: Some(false),
|
||||
output_format: Some("png".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let result = flux_1_dev_edit_image(request)
|
||||
.with_api_key(&args.api_key.0)
|
||||
.queue_webhook(args.webhook_url)
|
||||
.await;
|
||||
|
||||
result.map_err(|err| classify_fal_error(err))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::creds::fal_api_key::FalApiKey;
|
||||
use crate::requests::webhook::image::edit::enqueue_flux_1_dev_edit_image_webhook::{
|
||||
enqueue_flux_1_dev_edit_image_webhook, Flux1DevEditImageArgs, Flux1DevEditImageNumImages,
|
||||
};
|
||||
use errors::AnyhowResult;
|
||||
use std::fs::read_to_string;
|
||||
use test_data::web::image_urls::GHOST_IMAGE_URL;
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // manually run — fires a real Fal API request
|
||||
async fn test_single_image() -> AnyhowResult<()> {
|
||||
let secret = read_to_string("/Users/bt/Artcraft/credentials/fal_api_key.txt")?;
|
||||
let api_key = FalApiKey::from_str(&secret);
|
||||
|
||||
let args = Flux1DevEditImageArgs {
|
||||
prompt: "make this image look like a watercolor painting",
|
||||
image_url: GHOST_IMAGE_URL.to_string(),
|
||||
num_images: Flux1DevEditImageNumImages::One,
|
||||
api_key: &api_key,
|
||||
webhook_url: "https://example.com/webhook",
|
||||
};
|
||||
|
||||
let result = enqueue_flux_1_dev_edit_image_webhook(args).await?;
|
||||
assert!(result.request_id.is_some());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // manually run — fires a real Fal API request
|
||||
async fn test_batch_two() -> AnyhowResult<()> {
|
||||
let secret = read_to_string("/Users/bt/Artcraft/credentials/fal_api_key.txt")?;
|
||||
let api_key = FalApiKey::from_str(&secret);
|
||||
|
||||
let args = Flux1DevEditImageArgs {
|
||||
prompt: "turn this into a cyberpunk scene with neon lights",
|
||||
image_url: GHOST_IMAGE_URL.to_string(),
|
||||
num_images: Flux1DevEditImageNumImages::Two,
|
||||
api_key: &api_key,
|
||||
webhook_url: "https://example.com/webhook",
|
||||
};
|
||||
|
||||
let result = enqueue_flux_1_dev_edit_image_webhook(args).await?;
|
||||
assert!(result.request_id.is_some());
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
+120
@@ -0,0 +1,120 @@
|
||||
use crate::creds::fal_api_key::FalApiKey;
|
||||
use crate::error::classify_fal_error::classify_fal_error;
|
||||
use crate::error::fal_error_plus::FalErrorPlus;
|
||||
use crate::requests::http::image::edit::http_flux_1_schnell_edit_image::{flux_1_schnell_edit_image, Flux1SchnellEditImageInput};
|
||||
use crate::requests::api::webhook_response::WebhookResponse;
|
||||
use reqwest::IntoUrl;
|
||||
|
||||
pub struct Flux1SchnellEditImageArgs<'a, U: IntoUrl> {
|
||||
pub image_url: String,
|
||||
pub num_images: Flux1SchnellEditImageNumImages,
|
||||
pub image_size: Option<Flux1SchnellEditImageSize>,
|
||||
|
||||
// Fulfillment
|
||||
pub webhook_url: U,
|
||||
pub api_key: &'a FalApiKey,
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub enum Flux1SchnellEditImageNumImages {
|
||||
One,
|
||||
Two,
|
||||
Three,
|
||||
Four,
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub enum Flux1SchnellEditImageSize {
|
||||
Square,
|
||||
SquareHd,
|
||||
LandscapeFourByThree,
|
||||
LandscapeSixteenByNine,
|
||||
PortraitThreeByFour,
|
||||
PortraitNineBySixteen,
|
||||
}
|
||||
|
||||
pub async fn enqueue_flux_1_schnell_edit_image_webhook<U: IntoUrl>(
|
||||
args: Flux1SchnellEditImageArgs<'_, U>
|
||||
) -> Result<WebhookResponse, FalErrorPlus> {
|
||||
let num_images = match args.num_images {
|
||||
Flux1SchnellEditImageNumImages::One => 1,
|
||||
Flux1SchnellEditImageNumImages::Two => 2,
|
||||
Flux1SchnellEditImageNumImages::Three => 3,
|
||||
Flux1SchnellEditImageNumImages::Four => 4,
|
||||
};
|
||||
|
||||
let image_size = args.image_size.map(|s| match s {
|
||||
Flux1SchnellEditImageSize::Square => "square",
|
||||
Flux1SchnellEditImageSize::SquareHd => "square_hd",
|
||||
Flux1SchnellEditImageSize::LandscapeFourByThree => "landscape_4_3",
|
||||
Flux1SchnellEditImageSize::LandscapeSixteenByNine => "landscape_16_9",
|
||||
Flux1SchnellEditImageSize::PortraitThreeByFour => "portrait_4_3",
|
||||
Flux1SchnellEditImageSize::PortraitNineBySixteen => "portrait_16_9",
|
||||
}.to_string());
|
||||
|
||||
let request = Flux1SchnellEditImageInput {
|
||||
image_url: args.image_url,
|
||||
num_images: Some(num_images),
|
||||
image_size,
|
||||
enable_safety_checker: Some(false),
|
||||
output_format: Some("png".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let result = flux_1_schnell_edit_image(request)
|
||||
.with_api_key(&args.api_key.0)
|
||||
.queue_webhook(args.webhook_url)
|
||||
.await;
|
||||
|
||||
result.map_err(|err| classify_fal_error(err))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::creds::fal_api_key::FalApiKey;
|
||||
use crate::requests::webhook::image::edit::enqueue_flux_1_schnell_edit_image_webhook::{
|
||||
enqueue_flux_1_schnell_edit_image_webhook, Flux1SchnellEditImageArgs,
|
||||
Flux1SchnellEditImageNumImages, Flux1SchnellEditImageSize,
|
||||
};
|
||||
use errors::AnyhowResult;
|
||||
use std::fs::read_to_string;
|
||||
use test_data::web::image_urls::GHOST_IMAGE_URL;
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // manually run — fires a real Fal API request
|
||||
async fn test_single_image_no_size() -> AnyhowResult<()> {
|
||||
let secret = read_to_string("/Users/bt/Artcraft/credentials/fal_api_key.txt")?;
|
||||
let api_key = FalApiKey::from_str(&secret);
|
||||
|
||||
let args = Flux1SchnellEditImageArgs {
|
||||
image_url: GHOST_IMAGE_URL.to_string(),
|
||||
num_images: Flux1SchnellEditImageNumImages::One,
|
||||
image_size: None,
|
||||
api_key: &api_key,
|
||||
webhook_url: "https://example.com/webhook",
|
||||
};
|
||||
|
||||
let result = enqueue_flux_1_schnell_edit_image_webhook(args).await?;
|
||||
assert!(result.request_id.is_some());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // manually run — fires a real Fal API request
|
||||
async fn test_with_landscape_size() -> AnyhowResult<()> {
|
||||
let secret = read_to_string("/Users/bt/Artcraft/credentials/fal_api_key.txt")?;
|
||||
let api_key = FalApiKey::from_str(&secret);
|
||||
|
||||
let args = Flux1SchnellEditImageArgs {
|
||||
image_url: GHOST_IMAGE_URL.to_string(),
|
||||
num_images: Flux1SchnellEditImageNumImages::Two,
|
||||
image_size: Some(Flux1SchnellEditImageSize::LandscapeSixteenByNine),
|
||||
api_key: &api_key,
|
||||
webhook_url: "https://example.com/webhook",
|
||||
};
|
||||
|
||||
let result = enqueue_flux_1_schnell_edit_image_webhook(args).await?;
|
||||
assert!(result.request_id.is_some());
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,5 @@
|
||||
pub mod enqueue_flux_1_dev_edit_image_webhook;
|
||||
pub mod enqueue_flux_1_schnell_edit_image_webhook;
|
||||
pub mod enqueue_bytedance_seedream_v4_edit_image_webhook;
|
||||
pub mod enqueue_bytedance_seedream_v4p5_edit_image_webhook;
|
||||
pub mod enqueue_bytedance_seedream_v5_lite_edit_image_webhook;
|
||||
|
||||
+3
-3
@@ -23,10 +23,10 @@ pub async fn handle_image_edit_artcraft(
|
||||
match model {
|
||||
ImageEditModel::FluxProKontextMax => handle_artcraft_flux_kontext_edit(request, app, app_data_root, app_env_configs, storyteller_creds_manager).await,
|
||||
ImageEditModel::Gemini25Flash | ImageEditModel::NanoBanana => {
|
||||
handle_image_edit_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::NanaBanana, GenerationModel::NanoBanana).await
|
||||
handle_image_edit_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::NanoBanana, GenerationModel::NanoBanana).await
|
||||
},
|
||||
ImageEditModel::NanoBanana2 => handle_image_edit_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::NanaBanana2, GenerationModel::NanoBanana2).await,
|
||||
ImageEditModel::NanoBananaPro => handle_image_edit_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::NanaBananaPro, GenerationModel::NanoBananaPro).await,
|
||||
ImageEditModel::NanoBanana2 => handle_image_edit_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::NanoBanana2, GenerationModel::NanoBanana2).await,
|
||||
ImageEditModel::NanoBananaPro => handle_image_edit_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::NanoBananaPro, GenerationModel::NanoBananaPro).await,
|
||||
ImageEditModel::GptImage1 => handle_artcraft_gpt_image_1_edit(request, app, app_data_root, app_env_configs, storyteller_creds_manager).await,
|
||||
ImageEditModel::GptImage1p5 => handle_image_edit_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::GptImage1p5, GenerationModel::GptImage1p5).await,
|
||||
ImageEditModel::Seedream4 => handle_image_edit_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::Seedream4, GenerationModel::Seedream4).await,
|
||||
|
||||
+3
-3
@@ -38,10 +38,10 @@ pub async fn handle_text_to_image_artcraft(
|
||||
TextToImageModel::GptImage1 => handle_artcraft_gpt_image_1_text_to_image(request, app_env_configs, storyteller_creds_manager).await,
|
||||
TextToImageModel::GptImage1p5 => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::GptImage1p5, GenerationModel::GptImage1p5).await,
|
||||
TextToImageModel::Gemini25Flash | TextToImageModel::NanoBanana => {
|
||||
handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::NanaBanana, GenerationModel::NanoBanana).await
|
||||
handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::NanoBanana, GenerationModel::NanoBanana).await
|
||||
},
|
||||
TextToImageModel::NanoBanana2 => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::NanaBanana2, GenerationModel::NanoBanana2).await,
|
||||
TextToImageModel::NanoBananaPro => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::NanaBananaPro, GenerationModel::NanoBananaPro).await,
|
||||
TextToImageModel::NanoBanana2 => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::NanoBanana2, GenerationModel::NanoBanana2).await,
|
||||
TextToImageModel::NanoBananaPro => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::NanoBananaPro, GenerationModel::NanoBananaPro).await,
|
||||
TextToImageModel::Seedream4 => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::Seedream4, GenerationModel::Seedream4).await,
|
||||
TextToImageModel::Seedream4p5 => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::Seedream4p5, GenerationModel::Seedream4p5).await,
|
||||
TextToImageModel::Seedream5Lite => handle_text_to_image_artcraft_via_router(request, app_env_configs, storyteller_creds_manager, CommonImageModel::Seedream5Lite, GenerationModel::Seedream5Lite).await,
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
|
||||
/// The version of the ArtCraft desktop app reported in the UI and via Analytics.
|
||||
/// This should match the CI builds.
|
||||
pub const ARTCRAFT_VERSION: &str = "0.23.0";
|
||||
pub const ARTCRAFT_VERSION: &str = "0.24.0";
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"$schema": "https://schema.tauri.app/config/2.0.0-rc",
|
||||
"productName": "ArtCraft",
|
||||
"version": "0.23.0",
|
||||
"version": "0.24.0",
|
||||
"identifier": "ai.artcraft.app",
|
||||
"build": {
|
||||
"frontendDist": "../../../frontend/apps/artcraft/dist",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"$schema": "https://schema.tauri.app/config/2.0.0-rc",
|
||||
"productName": "ArtCraft",
|
||||
"version": "0.23.0",
|
||||
"version": "0.24.0",
|
||||
"identifier": "ai.artcraft.app",
|
||||
"build": {
|
||||
"frontendDist": "../../../frontend/apps/artcraft/dist",
|
||||
|
||||
@@ -6,12 +6,11 @@
|
||||
use actix_web::cookie::time::OffsetDateTime;
|
||||
use actix_web::cookie::{Cookie, SameSite};
|
||||
use actix_web::HttpRequest;
|
||||
use anyhow::anyhow;
|
||||
use log::warn;
|
||||
|
||||
use crate::sessions::http_user_session_payload::HttpUserSessionPayload;
|
||||
use crate::sessions::payload_signer::HttpUserSessionPayloadSigner;
|
||||
use errors::AnyhowResult;
|
||||
use crate::sessions::http_user_session_payload_signer::HttpUserSessionPayloadSigner;
|
||||
use crate::sessions::http_user_session_payload_error::HttpUserSessionPayloadError;
|
||||
use tokens::tokens::user_sessions::UserSessionToken;
|
||||
use tokens::tokens::users::UserToken;
|
||||
|
||||
@@ -31,19 +30,41 @@ pub struct HttpUserSessionManager {
|
||||
}
|
||||
|
||||
impl HttpUserSessionManager {
|
||||
pub fn new(cookie_domain: &str, hmac_secret: &str) -> AnyhowResult<Self> {
|
||||
pub fn new(cookie_domain: &str, hmac_secret: &str) -> Result<Self, HttpUserSessionPayloadError> {
|
||||
let payload_signer = HttpUserSessionPayloadSigner::new(hmac_secret)
|
||||
.map_err(|e| {
|
||||
warn!("Failed to construct HttpUserSessionPayloadSigner: {}", e);
|
||||
e
|
||||
})?;
|
||||
|
||||
Ok(Self {
|
||||
cookie_domain: cookie_domain.to_string(),
|
||||
payload_signer: HttpUserSessionPayloadSigner::new(hmac_secret)?,
|
||||
payload_signer,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn encode_session_payload(&self, session_token: &UserSessionToken, user_token: &UserToken) -> AnyhowResult<String> {
|
||||
pub fn encode_session_payload(
|
||||
&self,
|
||||
session_token: &UserSessionToken,
|
||||
user_token: &UserToken,
|
||||
) -> Result<String, HttpUserSessionPayloadError> {
|
||||
self.payload_signer.encode(session_token, user_token)
|
||||
.map_err(|e| {
|
||||
warn!("Failed to encode session payload: {}", e);
|
||||
e
|
||||
})
|
||||
}
|
||||
|
||||
pub fn create_cookie(&self, session_token: &UserSessionToken, user_token: &UserToken) -> AnyhowResult<Cookie> {
|
||||
let jwt_string = self.payload_signer.encode(session_token, user_token)?;
|
||||
pub fn create_cookie(
|
||||
&self,
|
||||
session_token: &UserSessionToken,
|
||||
user_token: &UserToken,
|
||||
) -> Result<Cookie, HttpUserSessionPayloadError> {
|
||||
let jwt_string = self.payload_signer.encode(session_token, user_token)
|
||||
.map_err(|e| {
|
||||
warn!("Failed to encode session cookie payload: {}", e);
|
||||
e
|
||||
})?;
|
||||
|
||||
let make_secure = !self.cookie_domain.to_lowercase().contains("jungle.horse")
|
||||
&& !self.cookie_domain.to_lowercase().contains("localhost");
|
||||
@@ -76,8 +97,8 @@ impl HttpUserSessionManager {
|
||||
|
||||
pub fn decode_session_payload_from_request(
|
||||
&self,
|
||||
request: &HttpRequest
|
||||
) -> AnyhowResult<Option<HttpUserSessionPayload>>
|
||||
request: &HttpRequest,
|
||||
) -> Result<Option<HttpUserSessionPayload>, HttpUserSessionPayloadError>
|
||||
{
|
||||
let signed_session_payload = self.session_payload_from_request(request)?;
|
||||
|
||||
@@ -88,8 +109,8 @@ impl HttpUserSessionManager {
|
||||
|
||||
match self.payload_signer.decode(&signed_session_payload) {
|
||||
Err(e) => {
|
||||
warn!("Session cookie decode error: {:?}", e);
|
||||
Err(anyhow!("Could not decode session cookie: {:?}", e))
|
||||
warn!("Session cookie decode error: {}", e);
|
||||
Err(e)
|
||||
},
|
||||
Ok(payload) => Ok(Some(payload)),
|
||||
}
|
||||
@@ -97,7 +118,10 @@ impl HttpUserSessionManager {
|
||||
|
||||
// NB: THIS IS ONLY FOR A QUICK HACK FOR FREAKING CORS UGH
|
||||
// THIS IS A HUGE STUPID SECURITY VULN. DAMNIT GOOGLE DAMNIT CORS.
|
||||
pub fn check_and_return_session_token_decodes(&self, request: &HttpRequest) -> AnyhowResult<Option<String>> {
|
||||
pub fn check_and_return_session_token_decodes(
|
||||
&self,
|
||||
request: &HttpRequest,
|
||||
) -> Result<Option<String>, HttpUserSessionPayloadError> {
|
||||
let signed_session_payload = self.session_payload_from_request(request)?;
|
||||
|
||||
let signed_session_payload = match signed_session_payload {
|
||||
@@ -107,8 +131,8 @@ impl HttpUserSessionManager {
|
||||
|
||||
match self.payload_signer.decode(&signed_session_payload) {
|
||||
Err(e) => {
|
||||
warn!("Session cookie decode error: {:?}", e);
|
||||
return Err(anyhow!("Could not decode session cookie: {:?}", e));
|
||||
warn!("Session cookie decode error: {}", e);
|
||||
return Err(e);
|
||||
},
|
||||
Ok(_payload) => {}, // Good! We'll just discard this.
|
||||
}
|
||||
@@ -116,15 +140,23 @@ impl HttpUserSessionManager {
|
||||
Ok(Some(signed_session_payload))
|
||||
}
|
||||
|
||||
fn session_payload_from_request(&self, request: &HttpRequest) -> AnyhowResult<Option<String>> {
|
||||
let mut signed_session_payload= request.cookie(SESSION_COOKIE_NAME)
|
||||
fn session_payload_from_request(
|
||||
&self,
|
||||
request: &HttpRequest,
|
||||
) -> Result<Option<String>, HttpUserSessionPayloadError> {
|
||||
let mut signed_session_payload = request.cookie(SESSION_COOKIE_NAME)
|
||||
.map(|cookie| cookie.value().to_string());
|
||||
|
||||
if signed_session_payload.is_none() {
|
||||
signed_session_payload = request.headers().get(SESSION_HEADER_NAME)
|
||||
.map(|header| header.to_str())
|
||||
.transpose()?
|
||||
.map(|payload| payload.to_string());
|
||||
let maybe_header = request.headers().get(SESSION_HEADER_NAME);
|
||||
if let Some(header) = maybe_header {
|
||||
let header_str = header.to_str()
|
||||
.map_err(|e| {
|
||||
warn!("Failed to read session HTTP header as str: {}", e);
|
||||
HttpUserSessionPayloadError::HttpSessionHeaderError(e.to_string())
|
||||
})?;
|
||||
signed_session_payload = Some(header_str.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(signed_session_payload)
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
use std::error::Error;
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
use jwt_signer::jwt_signer_error::JwtSignerError;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum HttpUserSessionPayloadError {
|
||||
/// Error reading HTTP header
|
||||
HttpSessionHeaderError(String),
|
||||
|
||||
/// Error encoding, decoding, or constructing the JWT signer.
|
||||
JwtSigner(JwtSignerError),
|
||||
}
|
||||
|
||||
impl Display for HttpUserSessionPayloadError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::HttpSessionHeaderError(msg) => write!(f, "HTTP session header error: {}", msg),
|
||||
Self::JwtSigner(e) => write!(f, "JWT signer error: {}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Error for HttpUserSessionPayloadError {}
|
||||
|
||||
impl From<JwtSignerError> for HttpUserSessionPayloadError {
|
||||
fn from(err: JwtSignerError) -> Self {
|
||||
Self::JwtSigner(err)
|
||||
}
|
||||
}
|
||||
+15
-9
@@ -1,8 +1,8 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::sessions::http_user_session_payload::HttpUserSessionPayload;
|
||||
use crate::sessions::http_user_session_payload_error::HttpUserSessionPayloadError;
|
||||
use jwt_signer::jwt_signer::JwtSigner;
|
||||
use errors::AnyhowResult;
|
||||
use tokens::tokens::user_sessions::UserSessionToken;
|
||||
use tokens::tokens::users::UserToken;
|
||||
|
||||
@@ -21,14 +21,18 @@ pub struct HttpUserSessionPayloadSigner {
|
||||
}
|
||||
|
||||
impl HttpUserSessionPayloadSigner {
|
||||
pub fn new(hmac_secret: &str) -> AnyhowResult<Self> {
|
||||
pub fn new(hmac_secret: &str) -> Result<Self, HttpUserSessionPayloadError> {
|
||||
Ok(Self {
|
||||
jwt_signer: JwtSigner::new(hmac_secret)?
|
||||
jwt_signer: JwtSigner::new(hmac_secret)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn encode(&self, session_token: &UserSessionToken, user_token: &UserToken) -> AnyhowResult<String> {
|
||||
let mut claims = BTreeMap::new();
|
||||
pub fn encode(
|
||||
&self,
|
||||
session_token: &UserSessionToken,
|
||||
user_token: &UserToken,
|
||||
) -> Result<String, HttpUserSessionPayloadError> {
|
||||
let mut claims: BTreeMap<&str, &str> = BTreeMap::new();
|
||||
let payload_version = PAYLOAD_VERSION.to_string();
|
||||
|
||||
claims.insert("session_token", session_token.as_str());
|
||||
@@ -36,12 +40,14 @@ impl HttpUserSessionPayloadSigner {
|
||||
claims.insert("version", &payload_version);
|
||||
|
||||
let jwt_string = self.jwt_signer.claims_to_jwt(&claims)?;
|
||||
|
||||
Ok(jwt_string)
|
||||
}
|
||||
|
||||
pub fn decode(&self, session_payload_contents: &str) -> AnyhowResult<HttpUserSessionPayload> {
|
||||
let claims = self.jwt_signer.jwt_to_claims(&session_payload_contents)?;
|
||||
pub fn decode(
|
||||
&self,
|
||||
session_payload_contents: &str,
|
||||
) -> Result<HttpUserSessionPayload, HttpUserSessionPayloadError> {
|
||||
let claims = self.jwt_signer.jwt_to_claims(session_payload_contents)?;
|
||||
|
||||
let session_token = claims["session_token"].clone();
|
||||
let maybe_user_token = claims.get("user_token")
|
||||
@@ -56,7 +62,7 @@ impl HttpUserSessionPayloadSigner {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::sessions::payload_signer::HttpUserSessionPayloadSigner;
|
||||
use crate::sessions::http_user_session_payload_signer::HttpUserSessionPayloadSigner;
|
||||
use tokens::tokens::user_sessions::UserSessionToken;
|
||||
use tokens::tokens::users::UserToken;
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod http_user_session_manager;
|
||||
pub mod http_user_session_payload;
|
||||
pub mod payload_signer;
|
||||
pub mod http_user_session_payload_error;
|
||||
pub mod http_user_session_payload_signer;
|
||||
@@ -19,9 +19,6 @@ normal = ["workspace-build-acceleration"]
|
||||
# Workspace hack (faster compile times)
|
||||
workspace-build-acceleration = { version = "0.1", path = "../../workspace-build-acceleration" }
|
||||
|
||||
# Internal
|
||||
errors.workspace = true
|
||||
|
||||
# External
|
||||
hmac = "0.10.1" # jwt signing
|
||||
jwt = "0.13.0"
|
||||
|
||||
@@ -6,7 +6,7 @@ use jwt::SignWithKey;
|
||||
use jwt::VerifyWithKey;
|
||||
use sha2::Sha256;
|
||||
|
||||
use errors::{anyhow, AnyhowResult};
|
||||
use crate::jwt_signer_error::JwtSignerError;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct JwtSigner {
|
||||
@@ -15,9 +15,9 @@ pub struct JwtSigner {
|
||||
|
||||
impl JwtSigner {
|
||||
|
||||
pub fn new(hmac_secret: &str) -> AnyhowResult<Self> {
|
||||
pub fn new(hmac_secret: &str) -> Result<Self, JwtSignerError> {
|
||||
let hmac_key = Hmac::new_varkey(hmac_secret.as_bytes())
|
||||
.map_err(|e| anyhow!("invalid hmac: {:?}", e))?;
|
||||
.map_err(|_| JwtSignerError::JwtInvalidKeyLength)?;
|
||||
|
||||
Ok(Self {
|
||||
hmac_key,
|
||||
@@ -25,33 +25,32 @@ impl JwtSigner {
|
||||
}
|
||||
|
||||
/// Turn a map of claims into a signed JWT cookie payload (string)
|
||||
pub fn claims_to_jwt(&self, claims: &impl Claimable) -> AnyhowResult<String> {
|
||||
let jwt_string = claims.sign(&self.hmac_key)?;
|
||||
Ok(jwt_string)
|
||||
pub fn claims_to_jwt(&self, claims: &impl Claimable) -> Result<String, JwtSignerError> {
|
||||
claims.sign(&self.hmac_key)
|
||||
}
|
||||
|
||||
/// Turn a JWT cookie payload (string) into a map of (key, value) claims
|
||||
pub fn jwt_to_claims(&self, jwt: &str) -> AnyhowResult<BTreeMap<String, String>> {
|
||||
let claims = jwt.verify_with_key(&self.hmac_key)?;
|
||||
Ok(claims)
|
||||
pub fn jwt_to_claims(&self, jwt: &str) -> Result<BTreeMap<String, String>, JwtSignerError> {
|
||||
jwt.verify_with_key(&self.hmac_key)
|
||||
.map_err(JwtSignerError::JwtVerifyError)
|
||||
}
|
||||
}
|
||||
|
||||
pub trait Claimable {
|
||||
fn sign(&self, hmac_key: &Hmac<Sha256>) -> AnyhowResult<String>;
|
||||
fn sign(&self, hmac_key: &Hmac<Sha256>) -> Result<String, JwtSignerError>;
|
||||
}
|
||||
|
||||
impl Claimable for BTreeMap<String, String> {
|
||||
fn sign(&self, hmac_key: &Hmac<Sha256>) -> AnyhowResult<String> {
|
||||
let claims = self.sign_with_key(hmac_key)?;
|
||||
Ok(claims)
|
||||
fn sign(&self, hmac_key: &Hmac<Sha256>) -> Result<String, JwtSignerError> {
|
||||
self.sign_with_key(hmac_key)
|
||||
.map_err(JwtSignerError::JwtSignError)
|
||||
}
|
||||
}
|
||||
|
||||
impl Claimable for BTreeMap<&str, &str> {
|
||||
fn sign(&self, hmac_key: &Hmac<Sha256>) -> AnyhowResult<String> {
|
||||
let claims = self.sign_with_key(hmac_key)?;
|
||||
Ok(claims)
|
||||
fn sign(&self, hmac_key: &Hmac<Sha256>) -> Result<String, JwtSignerError> {
|
||||
self.sign_with_key(hmac_key)
|
||||
.map_err(JwtSignerError::JwtSignError)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -92,4 +91,4 @@ mod tests {
|
||||
|
||||
assert_eq!("eyJhbGciOiJIUzI1NiJ9.eyJiYXIiOiJubyIsImZvbyI6InllcyJ9.k4eRvIrtw4tuGDtQqSmEDngtGYAaQVRYVMCK5N3JcZA", &jwt);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
use std::error::Error;
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum JwtSignerError {
|
||||
/// Error decoding and verifying a JWT to claims
|
||||
JwtVerifyError(jwt::Error),
|
||||
|
||||
/// Error encoding and signing JWT claims
|
||||
JwtSignError(jwt::Error),
|
||||
|
||||
/// Error constructing the JWT signer.
|
||||
JwtInvalidKeyLength,
|
||||
}
|
||||
|
||||
impl Display for JwtSignerError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::JwtVerifyError(e) => write!(f, "JWT verify error: {}", e),
|
||||
Self::JwtSignError(e) => write!(f, "JWT sign error: {}", e),
|
||||
Self::JwtInvalidKeyLength => write!(f, "JWT signer construction failed: invalid key length"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Error for JwtSignerError {}
|
||||
@@ -18,4 +18,5 @@
|
||||
#![allow(dead_code)]
|
||||
#![allow(non_snake_case)]
|
||||
|
||||
pub mod jwt_signer;
|
||||
pub mod jwt_signer;
|
||||
pub mod jwt_signer_error;
|
||||
+8
-14
@@ -2,13 +2,11 @@
|
||||
#![forbid(unused_imports)]
|
||||
#![forbid(unused_mut)]
|
||||
|
||||
use anyhow::anyhow;
|
||||
use log::warn;
|
||||
use sqlx::{Executor, MySql};
|
||||
use sqlx::pool::PoolConnection;
|
||||
|
||||
use enums::common::visibility::Visibility;
|
||||
use errors::AnyhowResult;
|
||||
use tokens::tokens::users::UserToken;
|
||||
|
||||
use crate::helpers::boolean_converters::{i8_to_bool, nullable_i8_to_bool_default_false, nullable_i8_to_optional_bool};
|
||||
@@ -99,7 +97,7 @@ impl SessionUserRecord {
|
||||
pub async fn get_user_session_by_token_pooled_connection(
|
||||
mysql_connection: &mut PoolConnection<MySql>,
|
||||
session_token: &str,
|
||||
) -> AnyhowResult<Option<SessionUserRecord>> {
|
||||
) -> Result<Option<SessionUserRecord>, sqlx::Error> {
|
||||
get_user_session_by_token(&mut **mysql_connection, session_token).await
|
||||
}
|
||||
|
||||
@@ -107,7 +105,7 @@ pub async fn get_user_session_by_token_pooled_connection(
|
||||
pub async fn get_user_session_by_token<'e, 'c : 'e, E>(
|
||||
mysql_executor: E,
|
||||
session_token: &str,
|
||||
) -> AnyhowResult<Option<SessionUserRecord>>
|
||||
) -> Result<Option<SessionUserRecord>, sqlx::Error>
|
||||
where E: 'e + Executor<'c, Database = MySql>
|
||||
{
|
||||
// NB: Lookup failure is Err(RowNotFound).
|
||||
@@ -236,17 +234,13 @@ WHERE user_sessions.token = ?
|
||||
|
||||
Ok(Some(result_user_record))
|
||||
},
|
||||
Err(sqlx::Error::RowNotFound) => {
|
||||
warn!("Valid cookie; invalid session: {}", session_token);
|
||||
Ok(None)
|
||||
},
|
||||
Err(err) => {
|
||||
match err {
|
||||
sqlx::Error::RowNotFound => {
|
||||
warn!("Valid cookie; invalid session: {}", session_token);
|
||||
Ok(None)
|
||||
},
|
||||
_ => {
|
||||
warn!("Session query error: {:?}", err);
|
||||
Err(anyhow!("session query error: {:?}", err))
|
||||
}
|
||||
}
|
||||
warn!("Session query error: {:?}", err);
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+7
-14
@@ -2,12 +2,9 @@
|
||||
#![forbid(unused_imports)]
|
||||
#![forbid(unused_mut)]
|
||||
|
||||
use anyhow::anyhow;
|
||||
use log::warn;
|
||||
use sqlx::{Executor, MySql};
|
||||
|
||||
use errors::AnyhowResult;
|
||||
|
||||
#[derive(Clone, Serialize, Deserialize)]
|
||||
pub struct SessionRecord {
|
||||
pub session_token: String,
|
||||
@@ -17,7 +14,7 @@ pub struct SessionRecord {
|
||||
pub async fn get_user_session_by_token_light<'e, 'c : 'e, E>(
|
||||
mysql_executor: E,
|
||||
session_token: &str,
|
||||
) -> AnyhowResult<Option<SessionRecord>>
|
||||
) -> Result<Option<SessionRecord>, sqlx::Error>
|
||||
where E: 'e + Executor<'c, Database = MySql>
|
||||
{
|
||||
let maybe_session_record = sqlx::query_as!(
|
||||
@@ -37,17 +34,13 @@ AND deleted_at IS NULL
|
||||
|
||||
match maybe_session_record {
|
||||
Ok(session_record) => Ok(Some(session_record)),
|
||||
Err(sqlx::Error::RowNotFound) => {
|
||||
warn!("Valid cookie; invalid session: {}", session_token);
|
||||
Ok(None)
|
||||
},
|
||||
Err(err) => {
|
||||
match err {
|
||||
sqlx::Error::RowNotFound => {
|
||||
warn!("Valid cookie; invalid session: {}", session_token);
|
||||
Ok(None)
|
||||
},
|
||||
_ => {
|
||||
warn!("Session query error: {:?}", err);
|
||||
Err(anyhow!("session query error: {:?}", err))
|
||||
}
|
||||
}
|
||||
warn!("Session query error: {:?}", err);
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,11 +16,11 @@ pub enum CommonImageModel {
|
||||
#[serde(rename = "gpt_image_1p5")]
|
||||
GptImage1p5,
|
||||
#[serde(rename = "nano_banana")]
|
||||
NanaBanana,
|
||||
NanoBanana,
|
||||
#[serde(rename = "nano_banana_2")]
|
||||
NanaBanana2,
|
||||
NanoBanana2,
|
||||
#[serde(rename = "nano_banana_pro")]
|
||||
NanaBananaPro,
|
||||
NanoBananaPro,
|
||||
#[serde(rename = "seedream_4")]
|
||||
Seedream4,
|
||||
#[serde(rename = "seedream_4p5")]
|
||||
@@ -28,3 +28,48 @@ pub enum CommonImageModel {
|
||||
#[serde(rename = "seedream_5_lite")]
|
||||
Seedream5Lite,
|
||||
}
|
||||
|
||||
impl CommonImageModel {
|
||||
pub fn to_common_model_type(&self) -> crate::common::generation::common_model_type::CommonModelType {
|
||||
use crate::common::generation::common_model_type::CommonModelType;
|
||||
match self {
|
||||
Self::Flux1Dev => CommonModelType::Flux1Dev,
|
||||
Self::Flux1Schnell => CommonModelType::Flux1Schnell,
|
||||
Self::FluxPro11 => CommonModelType::FluxPro11,
|
||||
Self::FluxPro11Ultra => CommonModelType::FluxPro11Ultra,
|
||||
Self::GptImage1p5 => CommonModelType::GptImage1p5,
|
||||
Self::NanoBanana => CommonModelType::NanoBanana,
|
||||
Self::NanoBanana2 => CommonModelType::NanoBanana2,
|
||||
Self::NanoBananaPro => CommonModelType::NanoBananaPro,
|
||||
Self::Seedream4 => CommonModelType::Seedream4,
|
||||
Self::Seedream4p5 => CommonModelType::Seedream4p5,
|
||||
Self::Seedream5Lite => CommonModelType::Seedream5Lite,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::common::generation::common_model_type::CommonModelType;
|
||||
|
||||
#[test]
|
||||
fn all_image_models_convert_to_common_model_type() {
|
||||
let models = [
|
||||
(CommonImageModel::Flux1Dev, CommonModelType::Flux1Dev),
|
||||
(CommonImageModel::Flux1Schnell, CommonModelType::Flux1Schnell),
|
||||
(CommonImageModel::FluxPro11, CommonModelType::FluxPro11),
|
||||
(CommonImageModel::FluxPro11Ultra, CommonModelType::FluxPro11Ultra),
|
||||
(CommonImageModel::GptImage1p5, CommonModelType::GptImage1p5),
|
||||
(CommonImageModel::NanoBanana, CommonModelType::NanoBanana),
|
||||
(CommonImageModel::NanoBanana2, CommonModelType::NanoBanana2),
|
||||
(CommonImageModel::NanoBananaPro, CommonModelType::NanoBananaPro),
|
||||
(CommonImageModel::Seedream4, CommonModelType::Seedream4),
|
||||
(CommonImageModel::Seedream4p5, CommonModelType::Seedream4p5),
|
||||
(CommonImageModel::Seedream5Lite, CommonModelType::Seedream5Lite),
|
||||
];
|
||||
for (image_model, expected) in models {
|
||||
assert_eq!(image_model.to_common_model_type(), expected);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -59,3 +59,62 @@ pub enum CommonVideoModel {
|
||||
#[serde(rename = "veo_3p1_fast")]
|
||||
Veo3p1Fast,
|
||||
}
|
||||
|
||||
impl CommonVideoModel {
|
||||
pub fn to_common_model_type(&self) -> crate::common::generation::common_model_type::CommonModelType {
|
||||
use crate::common::generation::common_model_type::CommonModelType;
|
||||
match self {
|
||||
Self::GrokVideo => CommonModelType::GrokVideo,
|
||||
Self::Kling16Pro => CommonModelType::Kling16Pro,
|
||||
Self::Kling21Pro => CommonModelType::Kling21Pro,
|
||||
Self::Kling21Master => CommonModelType::Kling21Master,
|
||||
Self::Kling2p5TurboPro => CommonModelType::Kling2p5TurboPro,
|
||||
Self::Kling2p6Pro => CommonModelType::Kling2p6Pro,
|
||||
Self::Kling3p0Standard => CommonModelType::Kling3p0Standard,
|
||||
Self::Kling3p0Pro => CommonModelType::Kling3p0Pro,
|
||||
Self::Seedance10Lite => CommonModelType::Seedance10Lite,
|
||||
Self::Seedance1p5Pro => CommonModelType::Seedance1p5Pro,
|
||||
Self::Seedance2p0 => CommonModelType::Seedance2p0,
|
||||
Self::Sora2 => CommonModelType::Sora2,
|
||||
Self::Sora2Pro => CommonModelType::Sora2Pro,
|
||||
Self::Veo2 => CommonModelType::Veo2,
|
||||
Self::Veo3 => CommonModelType::Veo3,
|
||||
Self::Veo3Fast => CommonModelType::Veo3Fast,
|
||||
Self::Veo3p1 => CommonModelType::Veo3p1,
|
||||
Self::Veo3p1Fast => CommonModelType::Veo3p1Fast,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::common::generation::common_model_type::CommonModelType;
|
||||
|
||||
#[test]
|
||||
fn all_video_models_convert_to_common_model_type() {
|
||||
let models = [
|
||||
(CommonVideoModel::GrokVideo, CommonModelType::GrokVideo),
|
||||
(CommonVideoModel::Kling16Pro, CommonModelType::Kling16Pro),
|
||||
(CommonVideoModel::Kling21Pro, CommonModelType::Kling21Pro),
|
||||
(CommonVideoModel::Kling21Master, CommonModelType::Kling21Master),
|
||||
(CommonVideoModel::Kling2p5TurboPro, CommonModelType::Kling2p5TurboPro),
|
||||
(CommonVideoModel::Kling2p6Pro, CommonModelType::Kling2p6Pro),
|
||||
(CommonVideoModel::Kling3p0Standard, CommonModelType::Kling3p0Standard),
|
||||
(CommonVideoModel::Kling3p0Pro, CommonModelType::Kling3p0Pro),
|
||||
(CommonVideoModel::Seedance10Lite, CommonModelType::Seedance10Lite),
|
||||
(CommonVideoModel::Seedance1p5Pro, CommonModelType::Seedance1p5Pro),
|
||||
(CommonVideoModel::Seedance2p0, CommonModelType::Seedance2p0),
|
||||
(CommonVideoModel::Sora2, CommonModelType::Sora2),
|
||||
(CommonVideoModel::Sora2Pro, CommonModelType::Sora2Pro),
|
||||
(CommonVideoModel::Veo2, CommonModelType::Veo2),
|
||||
(CommonVideoModel::Veo3, CommonModelType::Veo3),
|
||||
(CommonVideoModel::Veo3Fast, CommonModelType::Veo3Fast),
|
||||
(CommonVideoModel::Veo3p1, CommonModelType::Veo3p1),
|
||||
(CommonVideoModel::Veo3p1Fast, CommonModelType::Veo3p1Fast),
|
||||
];
|
||||
for (video_model, expected) in models {
|
||||
assert_eq!(video_model.to_common_model_type(), expected);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -59,7 +59,7 @@ fn build_omni_gen_image_models() -> Vec<OmniGenImageModelDetails> {
|
||||
// ...
|
||||
|
||||
models.push(OmniGenImageModelDetails {
|
||||
model: CommonImageModel::NanaBanana, // NB: currently Gemini25Flash in our system
|
||||
model: CommonImageModel::NanoBanana, // NB: currently Gemini25Flash in our system
|
||||
full_name: None,
|
||||
text_prompt_supported: Some(true),
|
||||
image_refs_supported: Some(true),
|
||||
@@ -85,7 +85,7 @@ fn build_omni_gen_image_models() -> Vec<OmniGenImageModelDetails> {
|
||||
});
|
||||
|
||||
models.push(OmniGenImageModelDetails {
|
||||
model: CommonImageModel::NanaBanana2,
|
||||
model: CommonImageModel::NanoBanana2,
|
||||
full_name: None,
|
||||
text_prompt_supported: Some(true),
|
||||
image_refs_supported: Some(true),
|
||||
@@ -118,7 +118,7 @@ fn build_omni_gen_image_models() -> Vec<OmniGenImageModelDetails> {
|
||||
});
|
||||
|
||||
models.push(OmniGenImageModelDetails {
|
||||
model: CommonImageModel::NanaBananaPro,
|
||||
model: CommonImageModel::NanoBananaPro,
|
||||
full_name: None,
|
||||
text_prompt_supported: Some(true),
|
||||
image_refs_supported: Some(true),
|
||||
@@ -167,15 +167,15 @@ fn build_omni_gen_image_model_providers() -> Vec<OmniGenImageModelProviderDetail
|
||||
overrides: None,
|
||||
},
|
||||
OmniGenImageProviderModelDetails {
|
||||
model: CommonImageModel::NanaBanana,
|
||||
model: CommonImageModel::NanoBanana,
|
||||
overrides: None,
|
||||
},
|
||||
OmniGenImageProviderModelDetails {
|
||||
model: CommonImageModel::NanaBanana2,
|
||||
model: CommonImageModel::NanoBanana2,
|
||||
overrides: None,
|
||||
},
|
||||
OmniGenImageProviderModelDetails {
|
||||
model: CommonImageModel::NanaBananaPro,
|
||||
model: CommonImageModel::NanoBananaPro,
|
||||
overrides: None,
|
||||
},
|
||||
],
|
||||
|
||||
@@ -905,7 +905,6 @@ use crate::http_server::endpoints::media_files::list::list_batch_generated_redux
|
||||
ListPinnedMediaFilesSuccessResponse,
|
||||
ListPinnedWeightsError,
|
||||
ListPinnedWeightsSuccessResponse,
|
||||
ListSessionJobsError,
|
||||
ListSessionJobsItem,
|
||||
ListSessionJobsQueryParams,
|
||||
ListSessionJobsSuccessResponse,
|
||||
@@ -970,7 +969,6 @@ use crate::http_server::endpoints::media_files::list::list_batch_generated_redux
|
||||
SearchModelWeightsSortDirection,
|
||||
SearchModelWeightsSortField,
|
||||
SearchModelWeightsSuccessResponse,
|
||||
SessionInfoError,
|
||||
SessionTokenInfoError,
|
||||
SessionTokenInfoSuccessResponse,
|
||||
SetMediaFileCoverImageError,
|
||||
|
||||
+28
-4
@@ -1,10 +1,13 @@
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::http_server::common_responses::common_web_error::CommonWebError;
|
||||
use crate::http_server::session::session_checker_error::SessionCheckerError;
|
||||
use crate::http_server::web_utils::user_session::require_user_session::RequireUserSessionError;
|
||||
use actix_artcraft::sessions::http_user_session_payload_error::HttpUserSessionPayloadError;
|
||||
use actix_http::StatusCode;
|
||||
use actix_web::{HttpResponse, HttpResponseBuilder, ResponseError};
|
||||
use anyhow::anyhow;
|
||||
use crate::http_server::web_utils::user_session::require_user_session::RequireUserSessionError;
|
||||
|
||||
/// An error type for actix-web handlers that wraps causal errors for debugging
|
||||
/// and paging while presenting safe, generic HTTP responses to users.
|
||||
@@ -194,9 +197,30 @@ impl From<RequireUserSessionError> for AdvancedCommonWebError {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<crate::http_server::common_responses::common_web_error::CommonWebError> for AdvancedCommonWebError {
|
||||
fn from(value: crate::http_server::common_responses::common_web_error::CommonWebError) -> Self {
|
||||
use crate::http_server::common_responses::common_web_error::CommonWebError;
|
||||
impl From<HttpUserSessionPayloadError> for AdvancedCommonWebError {
|
||||
fn from(_value: HttpUserSessionPayloadError) -> Self {
|
||||
// Any session error (bad header, bad JWT, construction failure) becomes 401.
|
||||
Self::NotAuthorized
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SessionCheckerError> for AdvancedCommonWebError {
|
||||
fn from(value: SessionCheckerError) -> Self {
|
||||
match value {
|
||||
// Bad / forged session cookie → 401 (don't page on this)
|
||||
// NOTE: If there's an elevated rate of across-the-board 401s,
|
||||
// then we probably misconfigured the HMAC secret
|
||||
SessionCheckerError::SessionPayload(_) => Self::NotAuthorized,
|
||||
// Underlying DB / cache errors → 500 with paging
|
||||
SessionCheckerError::Sqlx(err) => Self::from_error(err),
|
||||
// Likely Redis caching middleware
|
||||
SessionCheckerError::OtherError(err) => Self::from_anyhow_error(err),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<CommonWebError> for AdvancedCommonWebError {
|
||||
fn from(value: CommonWebError) -> Self {
|
||||
match value {
|
||||
CommonWebError::BadInputWithSimpleMessage(msg) => Self::BadInputWithSimpleMessage(msg),
|
||||
CommonWebError::NotAuthorized => Self::NotAuthorized,
|
||||
|
||||
+22
@@ -4,6 +4,7 @@ use log::{error, warn};
|
||||
use mysql_queries::errors::mysql_error::{MysqlCrateErrorSubtype, MysqlError};
|
||||
use std::error::Error;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use crate::http_server::session::session_checker_error::SessionCheckerError;
|
||||
use crate::http_server::web_utils::user_session::require_user_session::RequireUserSessionError;
|
||||
|
||||
#[derive(Debug, utoipa::ToSchema)]
|
||||
@@ -109,6 +110,27 @@ impl From<RequireUserSessionError> for CommonWebError {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SessionCheckerError> for CommonWebError {
|
||||
fn from(value: SessionCheckerError) -> Self {
|
||||
match value {
|
||||
// Bad / forged session cookie → 401
|
||||
// NOTE: If there's an elevated rate of across-the-board 401s,
|
||||
// then we probably misconfigured the HMAC secret
|
||||
SessionCheckerError::SessionPayload(_) => Self::NotAuthorized,
|
||||
// Underlying DB / cache errors → 500
|
||||
SessionCheckerError::Sqlx(err) => {
|
||||
error!("SessionChecker sqlx error: {:?}", err);
|
||||
Self::ServerError
|
||||
},
|
||||
// Likely Redis caching middleware
|
||||
SessionCheckerError::OtherError(err) => {
|
||||
error!("SessionChecker other error: {:?}", err);
|
||||
Self::ServerError
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
+1
-1
@@ -64,7 +64,7 @@ pub async fn create_character_handler(
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("Session checker error: {:?}", e);
|
||||
AdvancedCommonWebError::from_anyhow_error(e)
|
||||
AdvancedCommonWebError::from(e)
|
||||
})?;
|
||||
|
||||
let user_session = match maybe_user_session {
|
||||
|
||||
+3
-3
@@ -130,9 +130,9 @@ fn map_image_model(model: CommonImageModel) -> Result<RouterImageModel, HandlerE
|
||||
CommonImageModel::FluxPro11 => RouterImageModel::FluxPro11,
|
||||
CommonImageModel::FluxPro11Ultra => RouterImageModel::FluxPro11Ultra,
|
||||
CommonImageModel::GptImage1p5 => RouterImageModel::GptImage1p5,
|
||||
CommonImageModel::NanaBanana => RouterImageModel::NanaBanana,
|
||||
CommonImageModel::NanaBanana2 => RouterImageModel::NanaBanana2,
|
||||
CommonImageModel::NanaBananaPro => RouterImageModel::NanaBananaPro,
|
||||
CommonImageModel::NanoBanana => RouterImageModel::NanoBanana,
|
||||
CommonImageModel::NanoBanana2 => RouterImageModel::NanoBanana2,
|
||||
CommonImageModel::NanoBananaPro => RouterImageModel::NanoBananaPro,
|
||||
CommonImageModel::Seedream4 => RouterImageModel::Seedream4,
|
||||
CommonImageModel::Seedream4p5 => RouterImageModel::Seedream4p5,
|
||||
CommonImageModel::Seedream5Lite => RouterImageModel::Seedream5Lite,
|
||||
|
||||
+1
-1
@@ -97,7 +97,7 @@ pub async fn seedance_2p0_multi_function_video_gen_handler(
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("Session checker error: {:?}", e);
|
||||
AdvancedCommonWebError::from_anyhow_error(e)
|
||||
AdvancedCommonWebError::from(e)
|
||||
})?;
|
||||
|
||||
let user_session = match maybe_user_session {
|
||||
|
||||
+18
-80
@@ -1,7 +1,7 @@
|
||||
use std::collections::HashSet;
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
use crate::http_server::common_responses::media::media_domain::MediaDomain;
|
||||
use crate::http_server::common_responses::media::media_links_builder::MediaLinksBuilder;
|
||||
use crate::http_server::endpoints::inference_job::utils::estimates::estimate_job_progress::estimate_job_progress;
|
||||
@@ -10,12 +10,9 @@ use crate::http_server::endpoints::inference_job::utils::extractors::extract_liv
|
||||
use crate::http_server::endpoints::inference_job::utils::extractors::extract_polymorphic_inference_args::extract_polymorphic_inference_args;
|
||||
use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain;
|
||||
use crate::http_server::web_utils::filter_model_name::maybe_filter_model_name;
|
||||
use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
|
||||
use crate::state::server_state::ServerState;
|
||||
use actix_web::error::ResponseError;
|
||||
use actix_web::http::StatusCode;
|
||||
use actix_web::web::Json;
|
||||
use actix_web::{web, HttpRequest, HttpResponse};
|
||||
use actix_web::{web, HttpRequest};
|
||||
use actix_web_lab::extract::Query;
|
||||
use artcraft_api_defs::jobs::list_session_jobs::{ListSessionJobsItem, ListSessionJobsQueryParams, ListSessionJobsSuccessResponse, ListSessionRequestDetailsResponse, ListSessionResultDetailsResponse, ListSessionStatusDetailsResponse};
|
||||
use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
|
||||
@@ -34,7 +31,7 @@ use redis_common::redis_keys::RedisKeys;
|
||||
use server_environment::ServerEnvironment;
|
||||
use tokens::tokens::generic_inference_jobs::InferenceJobToken;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
use utoipa::{IntoParams, ToSchema};
|
||||
use utoipa::IntoParams;
|
||||
|
||||
/// For certain jobs or job classes (eg. non-premium), we kill the jobs if the user hasn't
|
||||
/// maintained a keepalive. This prevents wasted work when users who are unlikely to return
|
||||
@@ -42,40 +39,6 @@ use utoipa::{IntoParams, ToSchema};
|
||||
/// typically do not require keepalive.
|
||||
const JOB_KEEPALIVE_TTL_SECONDS : u64 = 60 * 3;
|
||||
|
||||
#[derive(Debug, ToSchema)]
|
||||
pub enum ListSessionJobsError {
|
||||
ServerError,
|
||||
NotFound,
|
||||
NotAuthorized,
|
||||
}
|
||||
|
||||
impl ResponseError for ListSessionJobsError {
|
||||
fn status_code(&self) -> StatusCode {
|
||||
match *self {
|
||||
ListSessionJobsError::ServerError => StatusCode::INTERNAL_SERVER_ERROR,
|
||||
ListSessionJobsError::NotFound => StatusCode::NOT_FOUND,
|
||||
ListSessionJobsError::NotAuthorized => StatusCode::UNAUTHORIZED,
|
||||
}
|
||||
}
|
||||
|
||||
fn error_response(&self) -> HttpResponse {
|
||||
let error_reason = match self {
|
||||
Self::ServerError => "server error".to_string(),
|
||||
Self::NotFound => "not found".to_string(),
|
||||
Self::NotAuthorized => "not authorized".to_string(),
|
||||
};
|
||||
|
||||
to_simple_json_error(&error_reason, self.status_code())
|
||||
}
|
||||
}
|
||||
|
||||
// NB: Not using derive_more::Display since Clion doesn't understand it.
|
||||
impl fmt::Display for ListSessionJobsError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{:?}", self)
|
||||
}
|
||||
}
|
||||
|
||||
/// List job statuses for jobs that are associated with the user's session.
|
||||
///
|
||||
/// The user must be logged in. This endpoint will only show the last 36 hours
|
||||
@@ -89,7 +52,8 @@ impl fmt::Display for ListSessionJobsError {
|
||||
path = "/v1/jobs/session",
|
||||
responses(
|
||||
(status = 200, body = ListSessionJobsSuccessResponse),
|
||||
(status = 500, body = ListSessionJobsError),
|
||||
(status = 401, description = "Not authorized"),
|
||||
(status = 500, description = "Server error"),
|
||||
),
|
||||
params(
|
||||
ListSessionJobsQueryParams
|
||||
@@ -99,29 +63,21 @@ pub async fn list_session_jobs_handler(
|
||||
http_request: HttpRequest,
|
||||
query: Query<ListSessionJobsQueryParams>,
|
||||
server_state: web::Data<Arc<ServerState>>
|
||||
) -> Result<Json<ListSessionJobsSuccessResponse>, ListSessionJobsError> {
|
||||
) -> Result<Json<ListSessionJobsSuccessResponse>, AdvancedCommonWebError> {
|
||||
|
||||
let mut mysql_connection = server_state.mysql_pool
|
||||
.acquire()
|
||||
.await
|
||||
.map_err(|err| {
|
||||
warn!("MySql pool error: {:?}", err);
|
||||
ListSessionJobsError::ServerError
|
||||
})?;
|
||||
let mut mysql_connection = server_state.mysql_pool.acquire().await?;
|
||||
|
||||
let maybe_avt_token = server_state.avt_cookie_manager
|
||||
.get_avt_token_from_request(&http_request);
|
||||
|
||||
// ==================== USER SESSION ==================== //
|
||||
|
||||
// NB: SessionCheckerError → AdvancedCommonWebError maps payload errors (bad cookie) to 401
|
||||
// and DB / cache errors to 500.
|
||||
let maybe_user_session = server_state
|
||||
.session_checker
|
||||
.maybe_get_user_session_extended_from_connection(&http_request, &mut mysql_connection)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("Session checker error: {:?}", e);
|
||||
ListSessionJobsError::ServerError
|
||||
})?;
|
||||
.await?;
|
||||
|
||||
let include_states = query.include_states
|
||||
.as_deref()
|
||||
@@ -135,20 +91,14 @@ pub async fn list_session_jobs_handler(
|
||||
.filter_map(|status| JobStatusPlus::from_str(status).ok())
|
||||
.collect::<HashSet<_>>());
|
||||
|
||||
let user;
|
||||
|
||||
match (maybe_user_session.as_ref(), maybe_avt_token.as_ref()) {
|
||||
(Some(session), _) => {
|
||||
user = SessionUser::User(&session.user_token_typed);
|
||||
}
|
||||
(None, Some(avt_token)) => {
|
||||
user = SessionUser::Anonymous(avt_token);
|
||||
}
|
||||
let user = match (maybe_user_session.as_ref(), maybe_avt_token.as_ref()) {
|
||||
(Some(session), _) => SessionUser::User(&session.user_token_typed),
|
||||
(None, Some(avt_token)) => SessionUser::Anonymous(avt_token),
|
||||
(None, None) => {
|
||||
// TODO(bt,2025-04-15): We should install an AVT cookie.
|
||||
return Err(ListSessionJobsError::NotAuthorized);
|
||||
return Err(AdvancedCommonWebError::NotAuthorized);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let args = ListSessionJobsForUserArgs {
|
||||
user,
|
||||
@@ -157,23 +107,11 @@ pub async fn list_session_jobs_handler(
|
||||
};
|
||||
|
||||
// NB: Since this is publicly exposed, we don't query sensitive data.
|
||||
let maybe_records = list_session_jobs_from_connection(
|
||||
args, &mut mysql_connection).await;
|
||||
|
||||
let records = match maybe_records {
|
||||
Ok(records) => records,
|
||||
Err(err) => {
|
||||
error!("tts job query error: {:?}", err);
|
||||
return Err(ListSessionJobsError::ServerError);
|
||||
}
|
||||
};
|
||||
let records = list_session_jobs_from_connection(args, &mut mysql_connection).await?;
|
||||
|
||||
let mut redis = server_state.redis_pool
|
||||
.get()
|
||||
.map_err(|e| {
|
||||
error!("redis error: {:?}", e);
|
||||
ListSessionJobsError::ServerError
|
||||
})?;
|
||||
.map_err(AdvancedCommonWebError::from_error)?;
|
||||
|
||||
// TODO(bt,2024-04-22): Look up the extra redis statuses per item.
|
||||
|
||||
@@ -204,7 +142,7 @@ fn records_to_response(
|
||||
records: Vec<GenericInferenceJobStatus>,
|
||||
server_environment: ServerEnvironment,
|
||||
media_domain: MediaDomain,
|
||||
) -> Result<Json<ListSessionJobsSuccessResponse>, ListSessionJobsError> {
|
||||
) -> Result<Json<ListSessionJobsSuccessResponse>, AdvancedCommonWebError> {
|
||||
let mut records = records.into_iter()
|
||||
.map(|record| {
|
||||
db_record_to_response_payload(record, None, server_environment, media_domain)
|
||||
|
||||
+26
-6
@@ -2,11 +2,13 @@ use std::sync::Arc;
|
||||
|
||||
use actix_web::web::Json;
|
||||
use actix_web::{web, HttpRequest};
|
||||
|
||||
use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest;
|
||||
use artcraft_api_defs::omni_gen::cost_response::omni_gen_image_cost_response::OmniGenImageCostResponse;
|
||||
use artcraft_router::api::provider::Provider;
|
||||
use log::warn;
|
||||
|
||||
use crate::http_server::common_responses::common_web_error::CommonWebError;
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
use crate::http_server::endpoints::omni_gen::generate::image::transform_request::transform_request;
|
||||
use crate::state::server_state::ServerState;
|
||||
|
||||
/// Estimate the cost of an image generation.
|
||||
@@ -18,14 +20,32 @@ use crate::state::server_state::ServerState;
|
||||
responses(
|
||||
(status = 200, description = "Success", body = OmniGenImageCostResponse),
|
||||
(status = 400, description = "Bad input"),
|
||||
(status = 401, description = "Unauthorized"),
|
||||
(status = 500, description = "Server error"),
|
||||
),
|
||||
)]
|
||||
pub async fn omni_gen_image_cost_handler(
|
||||
_http_request: HttpRequest,
|
||||
_request: Json<OmniGenImageCostAndGenerateRequest>,
|
||||
request: Json<OmniGenImageCostAndGenerateRequest>,
|
||||
_server_state: web::Data<Arc<ServerState>>,
|
||||
) -> Result<Json<OmniGenImageCostResponse>, CommonWebError> {
|
||||
todo!()
|
||||
) -> Result<Json<OmniGenImageCostResponse>, AdvancedCommonWebError> {
|
||||
let mut generate_request = transform_request(&request)?;
|
||||
generate_request.provider = Provider::Artcraft;
|
||||
|
||||
let plan = generate_request.build()
|
||||
.map_err(|e| {
|
||||
warn!("Failed to build cost plan: {}", e);
|
||||
AdvancedCommonWebError::from_error(e)
|
||||
})?;
|
||||
|
||||
let estimate = plan.estimate_costs();
|
||||
|
||||
Ok(Json(OmniGenImageCostResponse {
|
||||
success: true,
|
||||
cost_in_credits: estimate.cost_in_credits,
|
||||
cost_in_usd_cents: estimate.cost_in_usd_cents,
|
||||
is_free: estimate.is_free,
|
||||
is_unlimited: estimate.is_unlimited,
|
||||
is_rate_limited: estimate.is_rate_limited,
|
||||
has_watermark: estimate.has_watermark,
|
||||
}))
|
||||
}
|
||||
|
||||
+26
-6
@@ -2,11 +2,13 @@ use std::sync::Arc;
|
||||
|
||||
use actix_web::web::Json;
|
||||
use actix_web::{web, HttpRequest};
|
||||
|
||||
use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_video_cost_and_generate_request::OmniGenVideoCostAndGenerateRequest;
|
||||
use artcraft_api_defs::omni_gen::cost_response::omni_gen_video_cost_response::OmniGenVideoCostResponse;
|
||||
use artcraft_router::api::provider::Provider;
|
||||
use log::warn;
|
||||
|
||||
use crate::http_server::common_responses::common_web_error::CommonWebError;
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
use crate::http_server::endpoints::omni_gen::generate::video::transform_request::transform_request;
|
||||
use crate::state::server_state::ServerState;
|
||||
|
||||
/// Estimate the cost of a video generation.
|
||||
@@ -18,14 +20,32 @@ use crate::state::server_state::ServerState;
|
||||
responses(
|
||||
(status = 200, description = "Success", body = OmniGenVideoCostResponse),
|
||||
(status = 400, description = "Bad input"),
|
||||
(status = 401, description = "Unauthorized"),
|
||||
(status = 500, description = "Server error"),
|
||||
),
|
||||
)]
|
||||
pub async fn omni_gen_video_cost_handler(
|
||||
_http_request: HttpRequest,
|
||||
_request: Json<OmniGenVideoCostAndGenerateRequest>,
|
||||
request: Json<OmniGenVideoCostAndGenerateRequest>,
|
||||
_server_state: web::Data<Arc<ServerState>>,
|
||||
) -> Result<Json<OmniGenVideoCostResponse>, CommonWebError> {
|
||||
todo!()
|
||||
) -> Result<Json<OmniGenVideoCostResponse>, AdvancedCommonWebError> {
|
||||
let mut generate_request = transform_request(&request)?;
|
||||
generate_request.provider = Provider::Artcraft;
|
||||
|
||||
let plan = generate_request.build()
|
||||
.map_err(|e| {
|
||||
warn!("Failed to build cost plan: {}", e);
|
||||
AdvancedCommonWebError::from_error(e)
|
||||
})?;
|
||||
|
||||
let estimate = plan.estimate_costs();
|
||||
|
||||
Ok(Json(OmniGenVideoCostResponse {
|
||||
success: true,
|
||||
cost_in_credits: estimate.cost_in_credits,
|
||||
cost_in_usd_cents: estimate.cost_in_usd_cents,
|
||||
is_free: estimate.is_free,
|
||||
is_unlimited: estimate.is_unlimited,
|
||||
is_rate_limited: estimate.is_rate_limited,
|
||||
has_watermark: estimate.has_watermark,
|
||||
}))
|
||||
}
|
||||
|
||||
+4
@@ -1 +1,5 @@
|
||||
pub mod omni_gen_image_generate_handler;
|
||||
mod request_to_costs;
|
||||
mod request_to_plan;
|
||||
mod resolve_media_tokens;
|
||||
pub(crate) mod transform_request;
|
||||
|
||||
+256
-6
@@ -2,13 +2,41 @@ use std::sync::Arc;
|
||||
|
||||
use actix_web::web::Json;
|
||||
use actix_web::{web, HttpRequest};
|
||||
use log::{error, info, warn};
|
||||
use sqlx::Acquire;
|
||||
|
||||
use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest;
|
||||
use artcraft_api_defs::omni_gen::generate_response::omni_gen_image_generate_response::OmniGenImageGenerateResponse;
|
||||
use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptContextSemanticType;
|
||||
use enums::by_table::prompts::prompt_type::PromptType;
|
||||
use enums::common::generation::common_generation_mode::CommonGenerationMode;
|
||||
use enums::common::generation::common_model_type::CommonModelType;
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::common::visibility::Visibility;
|
||||
use http_server_common::request::get_request_ip::get_request_ip;
|
||||
use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue_with_apriori_job_token::{
|
||||
insert_generic_inference_job_for_fal_queue_with_apriori_job_token,
|
||||
InsertGenericInferenceForFalWithAprioriJobTokenArgs,
|
||||
};
|
||||
use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::FalCategory;
|
||||
use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
|
||||
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 tokens::tokens::generic_inference_jobs::InferenceJobToken;
|
||||
|
||||
use crate::http_server::common_responses::common_web_error::CommonWebError;
|
||||
use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error;
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test;
|
||||
use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format;
|
||||
use crate::state::server_state::ServerState;
|
||||
|
||||
use super::request_to_costs::request_to_costs;
|
||||
use super::request_to_plan::request_to_plan;
|
||||
use super::resolve_media_tokens::{resolve_media_tokens, apply_resolved_media};
|
||||
use super::transform_request::transform_request;
|
||||
|
||||
/// Generate an image using the omni-gen unified endpoint.
|
||||
#[utoipa::path(
|
||||
post,
|
||||
@@ -24,9 +52,231 @@ use crate::state::server_state::ServerState;
|
||||
),
|
||||
)]
|
||||
pub async fn omni_gen_image_generate_handler(
|
||||
_http_request: HttpRequest,
|
||||
_request: Json<OmniGenImageCostAndGenerateRequest>,
|
||||
_server_state: web::Data<Arc<ServerState>>,
|
||||
) -> Result<Json<OmniGenImageGenerateResponse>, CommonWebError> {
|
||||
todo!()
|
||||
http_request: HttpRequest,
|
||||
request: Json<OmniGenImageCostAndGenerateRequest>,
|
||||
server_state: web::Data<Arc<ServerState>>,
|
||||
) -> Result<Json<OmniGenImageGenerateResponse>, AdvancedCommonWebError> {
|
||||
|
||||
payments_error_test(&request.prompt.as_deref().unwrap_or(""))?;
|
||||
|
||||
// ==================== TRANSFORM REQUEST + PLAN ==================== //
|
||||
|
||||
let mut generate_request = transform_request(&request)?;
|
||||
|
||||
let maybe_prompt_model_type: Option<CommonModelType> = request.model
|
||||
.as_ref()
|
||||
.map(|m| m.to_common_model_type());
|
||||
|
||||
// ==================== COST ==================== //
|
||||
|
||||
let cost_estimate = request_to_costs(&generate_request)?;
|
||||
|
||||
// ==================== SESSION ==================== //
|
||||
|
||||
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);
|
||||
AdvancedCommonWebError::from(e)
|
||||
})?;
|
||||
|
||||
let user_token = match maybe_user_session.as_ref() {
|
||||
Some(session) => &session.user_token,
|
||||
None => return Err(AdvancedCommonWebError::NotAuthorized),
|
||||
};
|
||||
|
||||
let maybe_avt_token = server_state
|
||||
.avt_cookie_manager
|
||||
.get_avt_token_from_request(&http_request);
|
||||
|
||||
// ==================== IDEMPOTENCY ==================== //
|
||||
|
||||
let idempotency_token = request.idempotency_token.as_deref()
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
if let Err(reason) = validate_idempotency_token_format(&idempotency_token) {
|
||||
return Err(AdvancedCommonWebError::BadInputWithSimpleMessage(reason));
|
||||
}
|
||||
|
||||
insert_idempotency_token(&idempotency_token, &mut *mysql_connection)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
error!("Error inserting idempotency token: {:?}", err);
|
||||
AdvancedCommonWebError::BadInputWithSimpleMessage("repeated idempotency token".to_string())
|
||||
})?;
|
||||
|
||||
// ==================== RESOLVE MEDIA TOKENS ==================== //
|
||||
|
||||
let resolved_media = resolve_media_tokens(
|
||||
&request,
|
||||
&http_request,
|
||||
&mut mysql_connection,
|
||||
server_state.server_environment,
|
||||
).await?;
|
||||
|
||||
apply_resolved_media(&mut generate_request, &resolved_media);
|
||||
|
||||
// ==================== PLAN ==================== //
|
||||
|
||||
let plan = request_to_plan(&mut generate_request)?;
|
||||
|
||||
// ==================== BILLING ==================== //
|
||||
|
||||
let cost = cost_estimate.cost_in_credits.unwrap_or(0);
|
||||
|
||||
info!("Charging wallet: {} credits", cost);
|
||||
|
||||
let apriori_job_token = InferenceJobToken::generate();
|
||||
|
||||
if cost > 0 {
|
||||
attempt_wallet_deduction_else_common_web_error(
|
||||
user_token,
|
||||
Some(apriori_job_token.as_str()),
|
||||
cost,
|
||||
&mut mysql_connection,
|
||||
).await?;
|
||||
}
|
||||
|
||||
// ==================== EXECUTE GENERATION ==================== //
|
||||
|
||||
let fal_client = artcraft_router::client::router_fal_client::RouterFalClient::new(
|
||||
server_state.fal.api_key.clone(),
|
||||
server_state.fal.webhook_url.clone(),
|
||||
);
|
||||
|
||||
let router_client = artcraft_router::client::router_client::RouterClient::Fal(fal_client);
|
||||
|
||||
let generation_response = plan.generate_image(&router_client)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("Image generation failed: {:?}", e);
|
||||
AdvancedCommonWebError::from_error(e)
|
||||
})?;
|
||||
|
||||
let external_job_id = match &generation_response {
|
||||
artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse::Artcraft(p) => {
|
||||
p.inference_job_token.as_str().to_string()
|
||||
}
|
||||
artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse::Fal(p) => {
|
||||
p.request_id.clone().unwrap_or_default()
|
||||
}
|
||||
};
|
||||
|
||||
// ==================== DB TRANSACTION ==================== //
|
||||
|
||||
let ip_address = get_request_ip(&http_request);
|
||||
|
||||
let mut transaction = mysql_connection
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|err| {
|
||||
error!("Error starting MySQL transaction: {:?}", err);
|
||||
AdvancedCommonWebError::from_error(err)
|
||||
})?;
|
||||
|
||||
// -- Prompt --
|
||||
|
||||
let generation_mode = if request.image_media_tokens.is_some() {
|
||||
CommonGenerationMode::Edit
|
||||
} else {
|
||||
CommonGenerationMode::Text
|
||||
};
|
||||
|
||||
let prompt_result = insert_prompt(InsertPromptArgs {
|
||||
maybe_apriori_prompt_token: None,
|
||||
prompt_type: PromptType::ArtcraftApp,
|
||||
maybe_creator_user_token: Some(user_token),
|
||||
maybe_model_type: maybe_prompt_model_type,
|
||||
maybe_generation_provider: Some(GenerationProvider::Artcraft),
|
||||
maybe_positive_prompt: request.prompt.as_deref(),
|
||||
maybe_negative_prompt: None,
|
||||
maybe_other_args: None,
|
||||
maybe_generation_mode: Some(generation_mode),
|
||||
maybe_aspect_ratio: None,
|
||||
maybe_resolution: None,
|
||||
maybe_batch_count: request.image_batch_count.map(|c| c as u8),
|
||||
maybe_generate_audio: None,
|
||||
maybe_duration_seconds: 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
|
||||
}
|
||||
};
|
||||
|
||||
// -- Prompt context items --
|
||||
|
||||
if let Some(token) = prompt_token.as_ref() {
|
||||
let mut context_items = Vec::new();
|
||||
|
||||
if let Some(ref_tokens) = &request.image_media_tokens {
|
||||
for media_token in ref_tokens {
|
||||
context_items.push(PromptContextItem {
|
||||
media_token: media_token.clone(),
|
||||
context_semantic_type: PromptContextSemanticType::Imgsrc,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if !context_items.is_empty() {
|
||||
if let Err(err) = insert_batch_prompt_context_items(InsertBatchArgs {
|
||||
prompt_token: token.clone(),
|
||||
items: context_items,
|
||||
transaction: &mut transaction,
|
||||
}).await {
|
||||
warn!("Error inserting batch prompt context items: {:?}", err);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// -- Inference job --
|
||||
|
||||
let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token(
|
||||
InsertGenericInferenceForFalWithAprioriJobTokenArgs {
|
||||
apriori_job_token: &apriori_job_token,
|
||||
uuid_idempotency_token: &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: Some(user_token),
|
||||
maybe_avt_token: maybe_avt_token.as_ref(),
|
||||
creator_ip_address: &ip_address,
|
||||
creator_set_visibility: Visibility::Public,
|
||||
mysql_executor: &mut *transaction,
|
||||
starting_job_status_override: None,
|
||||
maybe_frontend_failure_category: None,
|
||||
maybe_failure_reason: None,
|
||||
phantom: Default::default(),
|
||||
}
|
||||
).await;
|
||||
|
||||
let job_token = match db_result {
|
||||
Ok(token) => token,
|
||||
Err(err) => {
|
||||
warn!("Error inserting inference job: {:?}", err);
|
||||
return Err(AdvancedCommonWebError::from_error(err));
|
||||
}
|
||||
};
|
||||
|
||||
transaction.commit().await.map_err(|err| {
|
||||
error!("Error committing transaction: {:?}", err);
|
||||
AdvancedCommonWebError::from_error(err)
|
||||
})?;
|
||||
|
||||
Ok(Json(OmniGenImageGenerateResponse {
|
||||
success: true,
|
||||
inference_job_token: job_token,
|
||||
}))
|
||||
}
|
||||
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
use artcraft_router::api::provider::Provider;
|
||||
use artcraft_router::generate::generate_image::generate_image_request::GenerateImageRequest;
|
||||
use artcraft_router::generate::generate_image::image_generation_cost_estimate::ImageGenerationCostEstimate;
|
||||
use log::warn;
|
||||
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
|
||||
/// Estimate costs for an image generation request.
|
||||
/// Always uses the Artcraft provider for costing regardless of the execution provider.
|
||||
pub fn request_to_costs(
|
||||
request: &GenerateImageRequest<'_>,
|
||||
) -> Result<ImageGenerationCostEstimate, AdvancedCommonWebError> {
|
||||
let mut cost_request = GenerateImageRequest {
|
||||
provider: Provider::Artcraft,
|
||||
..*request
|
||||
};
|
||||
|
||||
let plan = cost_request.build()
|
||||
.map_err(|e| {
|
||||
warn!("Failed to build cost plan: {}", e);
|
||||
AdvancedCommonWebError::from_error(e)
|
||||
})?;
|
||||
|
||||
Ok(plan.estimate_costs())
|
||||
}
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
use artcraft_router::api::provider::Provider;
|
||||
use artcraft_router::generate::generate_image::generate_image_request::GenerateImageRequest;
|
||||
use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan;
|
||||
use log::warn;
|
||||
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
|
||||
/// Build an image generation plan from a transformed request.
|
||||
/// Overrides the provider to Fal for execution (costing uses Artcraft separately).
|
||||
pub fn request_to_plan<'a>(
|
||||
request: &'a mut GenerateImageRequest<'a>,
|
||||
) -> Result<ImageGenerationPlan<'a>, AdvancedCommonWebError> {
|
||||
request.provider = Provider::Fal;
|
||||
|
||||
request.build()
|
||||
.map_err(|e| {
|
||||
warn!("Failed to build image generation plan: {}", e);
|
||||
AdvancedCommonWebError::from_error(e)
|
||||
})
|
||||
}
|
||||
+74
@@ -0,0 +1,74 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use actix_web::HttpRequest;
|
||||
use log::info;
|
||||
use sqlx::pool::PoolConnection;
|
||||
use sqlx::MySql;
|
||||
|
||||
use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest;
|
||||
use artcraft_router::api::image_list_ref::ImageListRef;
|
||||
use artcraft_router::generate::generate_image::generate_image_request::GenerateImageRequest;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
use crate::util::lookup::lookup_image_urls_as_map::lookup_image_urls_as_map;
|
||||
use server_environment::ServerEnvironment;
|
||||
|
||||
/// Resolved media URLs that can be referenced by a GenerateImageRequest.
|
||||
/// The owned Vecs here outlive the request so the GenerateImageRequest can borrow them.
|
||||
pub struct ResolvedImageMedia {
|
||||
pub url_map: HashMap<MediaFileToken, String>,
|
||||
pub image_input_urls: Vec<String>,
|
||||
}
|
||||
|
||||
/// Collect all media file tokens from the raw HTTP request, query them from the database,
|
||||
/// and return resolved CDN URLs.
|
||||
pub async fn resolve_media_tokens(
|
||||
omni_request: &OmniGenImageCostAndGenerateRequest,
|
||||
http_request: &HttpRequest,
|
||||
mysql_connection: &mut PoolConnection<MySql>,
|
||||
server_environment: ServerEnvironment,
|
||||
) -> Result<ResolvedImageMedia, AdvancedCommonWebError> {
|
||||
let mut all_tokens: Vec<MediaFileToken> = Vec::new();
|
||||
|
||||
if let Some(tokens) = &omni_request.image_media_tokens {
|
||||
all_tokens.extend(tokens.iter().cloned());
|
||||
}
|
||||
|
||||
if all_tokens.is_empty() {
|
||||
return Ok(ResolvedImageMedia {
|
||||
url_map: HashMap::new(),
|
||||
image_input_urls: Vec::new(),
|
||||
});
|
||||
}
|
||||
|
||||
info!("Resolving {} media file tokens to CDN URLs", all_tokens.len());
|
||||
|
||||
let url_map = lookup_image_urls_as_map(
|
||||
http_request,
|
||||
mysql_connection,
|
||||
server_environment,
|
||||
&all_tokens,
|
||||
).await?;
|
||||
|
||||
let image_input_urls = omni_request.image_media_tokens.as_ref()
|
||||
.map(|tokens| tokens.iter().filter_map(|t| url_map.get(t).cloned()).collect())
|
||||
.unwrap_or_default();
|
||||
|
||||
Ok(ResolvedImageMedia {
|
||||
url_map,
|
||||
image_input_urls,
|
||||
})
|
||||
}
|
||||
|
||||
/// Apply resolved media URLs to a GenerateImageRequest, replacing MediaFileToken refs with URL refs.
|
||||
pub fn apply_resolved_media<'a>(
|
||||
request: &mut GenerateImageRequest<'a>,
|
||||
resolved: &'a ResolvedImageMedia,
|
||||
) {
|
||||
if !resolved.image_input_urls.is_empty() {
|
||||
request.image_inputs = Some(ImageListRef::Urls(&resolved.image_input_urls));
|
||||
} else {
|
||||
request.image_inputs = None;
|
||||
}
|
||||
}
|
||||
+82
@@ -0,0 +1,82 @@
|
||||
use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest;
|
||||
use artcraft_router::api::common_aspect_ratio::CommonAspectRatio as CommonAspectRatioRouter;
|
||||
use artcraft_router::api::common_image_model::CommonImageModel as CommonImageModelRouter;
|
||||
use artcraft_router::api::common_resolution::CommonResolution as CommonResolutionRouter;
|
||||
use artcraft_router::api::image_list_ref::ImageListRef;
|
||||
use artcraft_router::api::provider::Provider;
|
||||
use artcraft_router::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
|
||||
use artcraft_router::generate::generate_image::generate_image_request::GenerateImageRequest;
|
||||
use enums::common::generation::common_aspect_ratio::CommonAspectRatio as CommonAspectRatioEnum;
|
||||
use enums::common::generation::common_image_model::CommonImageModel as CommonImageModelEnum;
|
||||
use enums::common::generation::common_resolution::CommonResolution as CommonResolutionEnum;
|
||||
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
|
||||
fn convert_model(
|
||||
model: &CommonImageModelEnum,
|
||||
) -> Result<CommonImageModelRouter, AdvancedCommonWebError> {
|
||||
let json = serde_json::to_string(model)?;
|
||||
serde_json::from_str(&json).map_err(|e| {
|
||||
AdvancedCommonWebError::BadInputWithSimpleMessage(
|
||||
format!("Unsupported image model: {}", e),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn convert_aspect_ratio(
|
||||
ar: &CommonAspectRatioEnum,
|
||||
) -> Result<CommonAspectRatioRouter, AdvancedCommonWebError> {
|
||||
let json = serde_json::to_string(ar)?;
|
||||
serde_json::from_str(&json).map_err(|e| {
|
||||
AdvancedCommonWebError::BadInputWithSimpleMessage(
|
||||
format!("Unsupported aspect ratio: {}", e),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn convert_resolution(
|
||||
res: &CommonResolutionEnum,
|
||||
) -> Result<CommonResolutionRouter, AdvancedCommonWebError> {
|
||||
let json = serde_json::to_string(res)?;
|
||||
serde_json::from_str(&json).map_err(|e| {
|
||||
AdvancedCommonWebError::BadInputWithSimpleMessage(
|
||||
format!("Unsupported resolution: {}", e),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn transform_request(
|
||||
request: &OmniGenImageCostAndGenerateRequest,
|
||||
) -> Result<GenerateImageRequest<'_>, AdvancedCommonWebError> {
|
||||
let api_model = request.model.as_ref()
|
||||
.ok_or_else(|| AdvancedCommonWebError::BadInputWithSimpleMessage(
|
||||
"model is required".to_string(),
|
||||
))?;
|
||||
|
||||
let model = convert_model(api_model)?;
|
||||
|
||||
let aspect_ratio = request.aspect_ratio.as_ref()
|
||||
.map(convert_aspect_ratio)
|
||||
.transpose()?;
|
||||
|
||||
let resolution = request.resolution.as_ref()
|
||||
.map(convert_resolution)
|
||||
.transpose()?;
|
||||
|
||||
Ok(GenerateImageRequest {
|
||||
model,
|
||||
provider: Provider::Artcraft,
|
||||
prompt: request.prompt.as_deref(),
|
||||
image_inputs: request.image_media_tokens.as_ref()
|
||||
.map(ImageListRef::MediaFileTokens),
|
||||
resolution,
|
||||
aspect_ratio,
|
||||
image_batch_count: request.image_batch_count,
|
||||
horizontal_angle: request.horizontal_angle,
|
||||
vertical_angle: request.vertical_angle,
|
||||
zoom: request.zoom,
|
||||
request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade,
|
||||
generation_mode_mismatch_strategy: None,
|
||||
idempotency_token: request.idempotency_token.as_deref(),
|
||||
})
|
||||
}
|
||||
+4
@@ -1 +1,5 @@
|
||||
pub mod omni_gen_video_generate_handler;
|
||||
mod request_to_costs;
|
||||
mod request_to_plan;
|
||||
mod resolve_media_tokens;
|
||||
pub(crate) mod transform_request;
|
||||
|
||||
+281
-6
@@ -2,13 +2,41 @@ use std::sync::Arc;
|
||||
|
||||
use actix_web::web::Json;
|
||||
use actix_web::{web, HttpRequest};
|
||||
use log::{error, info, warn};
|
||||
use sqlx::Acquire;
|
||||
|
||||
use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_video_cost_and_generate_request::OmniGenVideoCostAndGenerateRequest;
|
||||
use artcraft_api_defs::omni_gen::generate_response::omni_gen_video_generate_response::OmniGenVideoGenerateResponse;
|
||||
use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptContextSemanticType;
|
||||
use enums::by_table::prompts::prompt_type::PromptType;
|
||||
use enums::common::generation::common_generation_mode::CommonGenerationMode;
|
||||
use enums::common::generation::common_model_type::CommonModelType;
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::common::visibility::Visibility;
|
||||
use http_server_common::request::get_request_ip::get_request_ip;
|
||||
use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue_with_apriori_job_token::{
|
||||
insert_generic_inference_job_for_fal_queue_with_apriori_job_token,
|
||||
InsertGenericInferenceForFalWithAprioriJobTokenArgs,
|
||||
};
|
||||
use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::FalCategory;
|
||||
use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
|
||||
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 tokens::tokens::generic_inference_jobs::InferenceJobToken;
|
||||
|
||||
use crate::http_server::common_responses::common_web_error::CommonWebError;
|
||||
use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error;
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test;
|
||||
use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format;
|
||||
use crate::state::server_state::ServerState;
|
||||
|
||||
use super::request_to_costs::request_to_costs;
|
||||
use super::request_to_plan::request_to_plan;
|
||||
use super::resolve_media_tokens::{resolve_media_tokens, apply_resolved_media};
|
||||
use super::transform_request::transform_request;
|
||||
|
||||
/// Generate a video using the omni-gen unified endpoint.
|
||||
#[utoipa::path(
|
||||
post,
|
||||
@@ -24,9 +52,256 @@ use crate::state::server_state::ServerState;
|
||||
),
|
||||
)]
|
||||
pub async fn omni_gen_video_generate_handler(
|
||||
_http_request: HttpRequest,
|
||||
_request: Json<OmniGenVideoCostAndGenerateRequest>,
|
||||
_server_state: web::Data<Arc<ServerState>>,
|
||||
) -> Result<Json<OmniGenVideoGenerateResponse>, CommonWebError> {
|
||||
todo!()
|
||||
http_request: HttpRequest,
|
||||
request: Json<OmniGenVideoCostAndGenerateRequest>,
|
||||
server_state: web::Data<Arc<ServerState>>,
|
||||
) -> Result<Json<OmniGenVideoGenerateResponse>, AdvancedCommonWebError> {
|
||||
|
||||
payments_error_test(&request.prompt.as_deref().unwrap_or(""))?;
|
||||
|
||||
// ==================== TRANSFORM REQUEST + PLAN ==================== //
|
||||
|
||||
let mut generate_request = transform_request(&request)?;
|
||||
|
||||
let maybe_prompt_model_type: Option<CommonModelType> = request.model
|
||||
.as_ref()
|
||||
.map(|m| m.to_common_model_type());
|
||||
|
||||
// ==================== COST ==================== //
|
||||
|
||||
let cost_estimate = request_to_costs(&generate_request)?;
|
||||
|
||||
// ==================== SESSION ==================== //
|
||||
|
||||
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);
|
||||
AdvancedCommonWebError::from(e)
|
||||
})?;
|
||||
|
||||
let user_token = match maybe_user_session.as_ref() {
|
||||
Some(session) => &session.user_token,
|
||||
None => return Err(AdvancedCommonWebError::NotAuthorized),
|
||||
};
|
||||
|
||||
let maybe_avt_token = server_state
|
||||
.avt_cookie_manager
|
||||
.get_avt_token_from_request(&http_request);
|
||||
|
||||
// ==================== IDEMPOTENCY ==================== //
|
||||
|
||||
let idempotency_token = request.idempotency_token.as_deref()
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
if let Err(reason) = validate_idempotency_token_format(&idempotency_token) {
|
||||
return Err(AdvancedCommonWebError::BadInputWithSimpleMessage(reason));
|
||||
}
|
||||
|
||||
insert_idempotency_token(&idempotency_token, &mut *mysql_connection)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
error!("Error inserting idempotency token: {:?}", err);
|
||||
AdvancedCommonWebError::BadInputWithSimpleMessage("repeated idempotency token".to_string())
|
||||
})?;
|
||||
|
||||
// ==================== RESOLVE MEDIA TOKENS ==================== //
|
||||
|
||||
let resolved_media = resolve_media_tokens(
|
||||
&request,
|
||||
&http_request,
|
||||
&mut mysql_connection,
|
||||
server_state.server_environment,
|
||||
).await?;
|
||||
|
||||
apply_resolved_media(&mut generate_request, &resolved_media);
|
||||
|
||||
// ==================== PLAN ==================== //
|
||||
|
||||
let plan = request_to_plan(&mut generate_request)?;
|
||||
|
||||
// ==================== BILLING ==================== //
|
||||
|
||||
let cost = cost_estimate.cost_in_credits.unwrap_or(0);
|
||||
|
||||
info!("Charging wallet: {} credits", cost);
|
||||
|
||||
let apriori_job_token = InferenceJobToken::generate();
|
||||
|
||||
if cost > 0 {
|
||||
attempt_wallet_deduction_else_common_web_error(
|
||||
user_token,
|
||||
Some(apriori_job_token.as_str()),
|
||||
cost,
|
||||
&mut mysql_connection,
|
||||
).await?;
|
||||
}
|
||||
|
||||
// ==================== EXECUTE GENERATION ==================== //
|
||||
|
||||
let fal_client = artcraft_router::client::router_fal_client::RouterFalClient::new(
|
||||
server_state.fal.api_key.clone(),
|
||||
server_state.fal.webhook_url.clone(),
|
||||
);
|
||||
let router_client = artcraft_router::client::router_client::RouterClient::Fal(fal_client);
|
||||
|
||||
let generation_response = plan.generate_video(&router_client)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("Video generation failed: {:?}", e);
|
||||
AdvancedCommonWebError::from_error(e)
|
||||
})?;
|
||||
|
||||
let external_job_id = match &generation_response {
|
||||
artcraft_router::generate::generate_video::generate_video_response::GenerateVideoResponse::Artcraft(p) => {
|
||||
p.inference_job_token.as_str().to_string()
|
||||
}
|
||||
artcraft_router::generate::generate_video::generate_video_response::GenerateVideoResponse::Muapi(p) => {
|
||||
p.request_id.as_str().to_string()
|
||||
}
|
||||
artcraft_router::generate::generate_video::generate_video_response::GenerateVideoResponse::Seedance2Pro(p) => {
|
||||
p.order_id.clone()
|
||||
}
|
||||
};
|
||||
|
||||
// ==================== DB TRANSACTION ==================== //
|
||||
|
||||
let ip_address = get_request_ip(&http_request);
|
||||
|
||||
let mut transaction = mysql_connection
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|err| {
|
||||
error!("Error starting MySQL transaction: {:?}", err);
|
||||
AdvancedCommonWebError::from_error(err)
|
||||
})?;
|
||||
|
||||
// -- Prompt --
|
||||
|
||||
let generation_mode = if request.start_frame_image_media_token.is_some() {
|
||||
CommonGenerationMode::Keyframe
|
||||
} else {
|
||||
CommonGenerationMode::Text
|
||||
};
|
||||
|
||||
let prompt_result = insert_prompt(InsertPromptArgs {
|
||||
maybe_apriori_prompt_token: None,
|
||||
prompt_type: PromptType::ArtcraftApp,
|
||||
maybe_creator_user_token: Some(user_token),
|
||||
maybe_model_type: maybe_prompt_model_type,
|
||||
maybe_generation_provider: Some(GenerationProvider::Artcraft),
|
||||
maybe_positive_prompt: request.prompt.as_deref(),
|
||||
maybe_negative_prompt: request.negative_prompt.as_deref(),
|
||||
maybe_other_args: None,
|
||||
maybe_generation_mode: Some(generation_mode),
|
||||
maybe_aspect_ratio: None,
|
||||
maybe_resolution: None,
|
||||
maybe_batch_count: request.video_batch_count.map(|c| c as u8),
|
||||
maybe_generate_audio: request.generate_audio,
|
||||
maybe_duration_seconds: request.duration_seconds.map(|d| d as u32),
|
||||
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
|
||||
}
|
||||
};
|
||||
|
||||
// -- Prompt context items --
|
||||
|
||||
if let Some(token) = prompt_token.as_ref() {
|
||||
let mut context_items = Vec::new();
|
||||
|
||||
if let Some(media_token) = &request.start_frame_image_media_token {
|
||||
context_items.push(PromptContextItem {
|
||||
media_token: media_token.clone(),
|
||||
context_semantic_type: PromptContextSemanticType::VidStartFrame,
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(media_token) = &request.end_frame_image_media_token {
|
||||
context_items.push(PromptContextItem {
|
||||
media_token: media_token.clone(),
|
||||
context_semantic_type: PromptContextSemanticType::VidEndFrame,
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(ref_tokens) = &request.reference_image_media_tokens {
|
||||
for media_token in ref_tokens {
|
||||
context_items.push(PromptContextItem {
|
||||
media_token: media_token.clone(),
|
||||
context_semantic_type: PromptContextSemanticType::Imgref,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref_tokens) = &request.reference_video_media_tokens {
|
||||
for media_token in ref_tokens {
|
||||
context_items.push(PromptContextItem {
|
||||
media_token: media_token.clone(),
|
||||
context_semantic_type: PromptContextSemanticType::VidRef,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if !context_items.is_empty() {
|
||||
if let Err(err) = insert_batch_prompt_context_items(InsertBatchArgs {
|
||||
prompt_token: token.clone(),
|
||||
items: context_items,
|
||||
transaction: &mut transaction,
|
||||
}).await {
|
||||
warn!("Error inserting batch prompt context items: {:?}", err);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// -- Inference job --
|
||||
|
||||
let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token(
|
||||
InsertGenericInferenceForFalWithAprioriJobTokenArgs {
|
||||
apriori_job_token: &apriori_job_token,
|
||||
uuid_idempotency_token: &idempotency_token,
|
||||
maybe_external_third_party_id: &external_job_id,
|
||||
fal_category: FalCategory::VideoGeneration,
|
||||
maybe_inference_args: None,
|
||||
maybe_prompt_token: prompt_token.as_ref(),
|
||||
maybe_creator_user_token: Some(user_token),
|
||||
maybe_avt_token: maybe_avt_token.as_ref(),
|
||||
creator_ip_address: &ip_address,
|
||||
creator_set_visibility: Visibility::Public,
|
||||
mysql_executor: &mut *transaction,
|
||||
starting_job_status_override: None,
|
||||
maybe_frontend_failure_category: None,
|
||||
maybe_failure_reason: None,
|
||||
phantom: Default::default(),
|
||||
}
|
||||
).await;
|
||||
|
||||
let job_token = match db_result {
|
||||
Ok(token) => token,
|
||||
Err(err) => {
|
||||
warn!("Error inserting inference job: {:?}", err);
|
||||
return Err(AdvancedCommonWebError::from_error(err));
|
||||
}
|
||||
};
|
||||
|
||||
transaction.commit().await.map_err(|err| {
|
||||
error!("Error committing transaction: {:?}", err);
|
||||
AdvancedCommonWebError::from_error(err)
|
||||
})?;
|
||||
|
||||
Ok(Json(OmniGenVideoGenerateResponse {
|
||||
success: true,
|
||||
inference_job_token: job_token,
|
||||
}))
|
||||
}
|
||||
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
use artcraft_router::api::provider::Provider;
|
||||
use artcraft_router::generate::generate_video::generate_video_request::GenerateVideoRequest;
|
||||
use artcraft_router::generate::generate_video::video_generation_cost_estimate::VideoGenerationCostEstimate;
|
||||
use log::warn;
|
||||
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
|
||||
/// Estimate costs for a video generation request.
|
||||
/// Always uses the Artcraft provider for costing regardless of the execution provider.
|
||||
pub fn request_to_costs(
|
||||
request: &GenerateVideoRequest<'_>,
|
||||
) -> Result<VideoGenerationCostEstimate, AdvancedCommonWebError> {
|
||||
let mut cost_request = GenerateVideoRequest {
|
||||
provider: Provider::Artcraft,
|
||||
..*request
|
||||
};
|
||||
|
||||
let plan = cost_request.build()
|
||||
.map_err(|e| {
|
||||
warn!("Failed to build cost plan: {}", e);
|
||||
AdvancedCommonWebError::from_error(e)
|
||||
})?;
|
||||
|
||||
Ok(plan.estimate_costs())
|
||||
}
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
use artcraft_router::api::provider::Provider;
|
||||
use artcraft_router::generate::generate_video::generate_video_request::GenerateVideoRequest;
|
||||
use artcraft_router::generate::generate_video::video_generation_plan::VideoGenerationPlan;
|
||||
use log::warn;
|
||||
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
|
||||
/// Build a video generation plan from a transformed request.
|
||||
/// Overrides the provider to Fal for execution (costing uses Artcraft separately).
|
||||
pub fn request_to_plan<'a>(
|
||||
request: &'a mut GenerateVideoRequest<'a>,
|
||||
) -> Result<VideoGenerationPlan<'a>, AdvancedCommonWebError> {
|
||||
request.provider = Provider::Fal;
|
||||
|
||||
request.build()
|
||||
.map_err(|e| {
|
||||
warn!("Failed to build video generation plan: {}", e);
|
||||
AdvancedCommonWebError::from_error(e)
|
||||
})
|
||||
}
|
||||
+133
@@ -0,0 +1,133 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use actix_web::HttpRequest;
|
||||
use log::info;
|
||||
use sqlx::pool::PoolConnection;
|
||||
use sqlx::MySql;
|
||||
|
||||
use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_video_cost_and_generate_request::OmniGenVideoCostAndGenerateRequest;
|
||||
use artcraft_router::api::audio_list_ref::AudioListRef;
|
||||
use artcraft_router::api::image_list_ref::ImageListRef;
|
||||
use artcraft_router::api::image_ref::ImageRef;
|
||||
use artcraft_router::api::video_list_ref::VideoListRef;
|
||||
use artcraft_router::generate::generate_video::generate_video_request::GenerateVideoRequest;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
use crate::util::lookup::lookup_image_urls_as_map::lookup_image_urls_as_map;
|
||||
use server_environment::ServerEnvironment;
|
||||
|
||||
/// Resolved media URLs that can be referenced by a GenerateVideoRequest.
|
||||
/// The owned Vecs here outlive the request so the GenerateVideoRequest can borrow them.
|
||||
pub struct ResolvedVideoMedia {
|
||||
pub url_map: HashMap<MediaFileToken, String>,
|
||||
pub maybe_start_frame_url: Option<String>,
|
||||
pub maybe_end_frame_url: Option<String>,
|
||||
pub reference_image_urls: Vec<String>,
|
||||
pub reference_video_urls: Vec<String>,
|
||||
pub reference_audio_urls: Vec<String>,
|
||||
}
|
||||
|
||||
/// Collect all media file tokens from the raw HTTP request, query them from the database,
|
||||
/// and return resolved CDN URLs.
|
||||
pub async fn resolve_media_tokens(
|
||||
omni_request: &OmniGenVideoCostAndGenerateRequest,
|
||||
http_request: &HttpRequest,
|
||||
mysql_connection: &mut PoolConnection<MySql>,
|
||||
server_environment: ServerEnvironment,
|
||||
) -> Result<ResolvedVideoMedia, AdvancedCommonWebError> {
|
||||
let mut all_tokens: Vec<MediaFileToken> = Vec::new();
|
||||
|
||||
if let Some(token) = &omni_request.start_frame_image_media_token {
|
||||
all_tokens.push(token.clone());
|
||||
}
|
||||
if let Some(token) = &omni_request.end_frame_image_media_token {
|
||||
all_tokens.push(token.clone());
|
||||
}
|
||||
if let Some(tokens) = &omni_request.reference_image_media_tokens {
|
||||
all_tokens.extend(tokens.iter().cloned());
|
||||
}
|
||||
if let Some(tokens) = &omni_request.reference_video_media_tokens {
|
||||
all_tokens.extend(tokens.iter().cloned());
|
||||
}
|
||||
if let Some(tokens) = &omni_request.reference_audio_media_tokens {
|
||||
all_tokens.extend(tokens.iter().cloned());
|
||||
}
|
||||
|
||||
if all_tokens.is_empty() {
|
||||
return Ok(ResolvedVideoMedia {
|
||||
url_map: HashMap::new(),
|
||||
maybe_start_frame_url: None,
|
||||
maybe_end_frame_url: None,
|
||||
reference_image_urls: Vec::new(),
|
||||
reference_video_urls: Vec::new(),
|
||||
reference_audio_urls: Vec::new(),
|
||||
});
|
||||
}
|
||||
|
||||
info!("Resolving {} media file tokens to CDN URLs", all_tokens.len());
|
||||
|
||||
let url_map = lookup_image_urls_as_map(
|
||||
http_request,
|
||||
mysql_connection,
|
||||
server_environment,
|
||||
&all_tokens,
|
||||
).await?;
|
||||
|
||||
let maybe_start_frame_url = omni_request.start_frame_image_media_token.as_ref()
|
||||
.and_then(|t| url_map.get(t).cloned());
|
||||
|
||||
let maybe_end_frame_url = omni_request.end_frame_image_media_token.as_ref()
|
||||
.and_then(|t| url_map.get(t).cloned());
|
||||
|
||||
let reference_image_urls = omni_request.reference_image_media_tokens.as_ref()
|
||||
.map(|tokens| tokens.iter().filter_map(|t| url_map.get(t).cloned()).collect())
|
||||
.unwrap_or_default();
|
||||
|
||||
let reference_video_urls = omni_request.reference_video_media_tokens.as_ref()
|
||||
.map(|tokens| tokens.iter().filter_map(|t| url_map.get(t).cloned()).collect())
|
||||
.unwrap_or_default();
|
||||
|
||||
let reference_audio_urls = omni_request.reference_audio_media_tokens.as_ref()
|
||||
.map(|tokens| tokens.iter().filter_map(|t| url_map.get(t).cloned()).collect())
|
||||
.unwrap_or_default();
|
||||
|
||||
Ok(ResolvedVideoMedia {
|
||||
url_map,
|
||||
maybe_start_frame_url,
|
||||
maybe_end_frame_url,
|
||||
reference_image_urls,
|
||||
reference_video_urls,
|
||||
reference_audio_urls,
|
||||
})
|
||||
}
|
||||
|
||||
/// Apply resolved media URLs to a GenerateVideoRequest, replacing MediaFileToken refs with URL refs.
|
||||
pub fn apply_resolved_media<'a>(
|
||||
request: &mut GenerateVideoRequest<'a>,
|
||||
resolved: &'a ResolvedVideoMedia,
|
||||
) {
|
||||
request.start_frame = resolved.maybe_start_frame_url.as_deref()
|
||||
.map(ImageRef::Url);
|
||||
|
||||
request.end_frame = resolved.maybe_end_frame_url.as_deref()
|
||||
.map(ImageRef::Url);
|
||||
|
||||
if !resolved.reference_image_urls.is_empty() {
|
||||
request.reference_images = Some(ImageListRef::Urls(&resolved.reference_image_urls));
|
||||
} else {
|
||||
request.reference_images = None;
|
||||
}
|
||||
|
||||
if !resolved.reference_video_urls.is_empty() {
|
||||
request.reference_videos = Some(VideoListRef::Urls(&resolved.reference_video_urls));
|
||||
} else {
|
||||
request.reference_videos = None;
|
||||
}
|
||||
|
||||
if !resolved.reference_audio_urls.is_empty() {
|
||||
request.reference_audio = Some(AudioListRef::Urls(&resolved.reference_audio_urls));
|
||||
} else {
|
||||
request.reference_audio = None;
|
||||
}
|
||||
}
|
||||
+92
@@ -0,0 +1,92 @@
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_video_cost_and_generate_request::OmniGenVideoCostAndGenerateRequest;
|
||||
use artcraft_router::api::audio_list_ref::AudioListRef;
|
||||
use artcraft_router::api::common_aspect_ratio::CommonAspectRatio as CommonAspectRatioRouter;
|
||||
use artcraft_router::api::common_resolution::CommonResolution as CommonResolutionRouter;
|
||||
use artcraft_router::api::common_video_model::CommonVideoModel as CommonVideoModelRouter;
|
||||
use artcraft_router::api::image_list_ref::ImageListRef;
|
||||
use artcraft_router::api::image_ref::ImageRef;
|
||||
use artcraft_router::api::provider::Provider;
|
||||
use artcraft_router::api::video_list_ref::VideoListRef;
|
||||
use artcraft_router::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
|
||||
use artcraft_router::generate::generate_video::generate_video_request::GenerateVideoRequest;
|
||||
use enums::common::generation::common_aspect_ratio::CommonAspectRatio as CommonAspectRatioEnum;
|
||||
use enums::common::generation::common_resolution::CommonResolution as CommonResolutionEnum;
|
||||
use enums::common::generation::common_video_model::CommonVideoModel as CommonVideoModelEnum;
|
||||
|
||||
fn convert_model(
|
||||
model: &CommonVideoModelEnum,
|
||||
) -> Result<CommonVideoModelRouter, AdvancedCommonWebError> {
|
||||
let json = serde_json::to_string(model)?;
|
||||
serde_json::from_str(&json).map_err(|e| {
|
||||
AdvancedCommonWebError::BadInputWithSimpleMessage(
|
||||
format!("Unsupported video model: {}", e),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn convert_aspect_ratio(
|
||||
ar: &CommonAspectRatioEnum,
|
||||
) -> Result<CommonAspectRatioRouter, AdvancedCommonWebError> {
|
||||
let json = serde_json::to_string(ar)?;
|
||||
serde_json::from_str(&json).map_err(|e| {
|
||||
AdvancedCommonWebError::BadInputWithSimpleMessage(
|
||||
format!("Unsupported aspect ratio: {}", e),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn convert_resolution(
|
||||
res: &CommonResolutionEnum,
|
||||
) -> Result<CommonResolutionRouter, AdvancedCommonWebError> {
|
||||
let json = serde_json::to_string(res)?;
|
||||
serde_json::from_str(&json).map_err(|e| {
|
||||
AdvancedCommonWebError::BadInputWithSimpleMessage(
|
||||
format!("Unsupported resolution: {}", e),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn transform_request(
|
||||
request: &OmniGenVideoCostAndGenerateRequest,
|
||||
) -> Result<GenerateVideoRequest<'_>, AdvancedCommonWebError> {
|
||||
let api_model = request.model.as_ref()
|
||||
.ok_or_else(|| AdvancedCommonWebError::BadInputWithSimpleMessage(
|
||||
"model is required".to_string(),
|
||||
))?;
|
||||
|
||||
let model = convert_model(api_model)?;
|
||||
|
||||
let aspect_ratio = request.aspect_ratio.as_ref()
|
||||
.map(convert_aspect_ratio)
|
||||
.transpose()?;
|
||||
|
||||
let resolution = request.resolution.as_ref()
|
||||
.map(convert_resolution)
|
||||
.transpose()?;
|
||||
|
||||
Ok(GenerateVideoRequest {
|
||||
model,
|
||||
provider: Provider::Artcraft,
|
||||
prompt: request.prompt.as_deref(),
|
||||
negative_prompt: request.negative_prompt.as_deref(),
|
||||
start_frame: request.start_frame_image_media_token.as_ref()
|
||||
.map(ImageRef::MediaFileToken),
|
||||
end_frame: request.end_frame_image_media_token.as_ref()
|
||||
.map(ImageRef::MediaFileToken),
|
||||
reference_images: request.reference_image_media_tokens.as_ref()
|
||||
.map(ImageListRef::MediaFileTokens),
|
||||
reference_videos: request.reference_video_media_tokens.as_ref()
|
||||
.map(VideoListRef::MediaFileTokens),
|
||||
reference_audio: request.reference_audio_media_tokens.as_ref()
|
||||
.map(AudioListRef::MediaFileTokens),
|
||||
reference_character_tokens: None,
|
||||
resolution,
|
||||
aspect_ratio,
|
||||
duration_seconds: request.duration_seconds,
|
||||
video_batch_count: request.video_batch_count,
|
||||
generate_audio: request.generate_audio,
|
||||
request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade,
|
||||
idempotency_token: request.idempotency_token.as_deref(),
|
||||
})
|
||||
}
|
||||
+61
-111
@@ -3,60 +3,26 @@
|
||||
#![forbid(unused_mut)]
|
||||
#![forbid(unused_variables)]
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use actix_web::error::ResponseError;
|
||||
use actix_web::http::StatusCode;
|
||||
use actix_web::{web, HttpRequest, HttpResponse};
|
||||
use log::warn;
|
||||
use sqlx::MySqlPool;
|
||||
use utoipa::ToSchema;
|
||||
|
||||
use artcraft_api_defs::users::session_info::{FakeYouPlan, SessionInfoSuccessResponse, SessionOnboardingState, SessionUserInfo, StorytellerStreamPlan};
|
||||
use enums::by_table::users::user_feature_flag::UserFeatureFlag;
|
||||
use http_server_common::response::response_error_helpers::to_simple_json_error;
|
||||
|
||||
use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
|
||||
use crate::http_server::common_responses::user_details_lite_builder::UserDetailsLightBuilder;
|
||||
use crate::http_server::cookies::anonymous_visitor_tracking::avt_cookie_manager::AvtCookieManager;
|
||||
use crate::http_server::session::lookup::user_session_feature_flags::UserSessionFeatureFlags;
|
||||
use crate::http_server::session::session_checker::SessionChecker;
|
||||
|
||||
|
||||
#[derive(Debug, ToSchema)]
|
||||
pub enum SessionInfoError {
|
||||
ServerError,
|
||||
}
|
||||
|
||||
impl ResponseError for SessionInfoError {
|
||||
fn status_code(&self) -> StatusCode {
|
||||
match *self {
|
||||
SessionInfoError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR,
|
||||
}
|
||||
}
|
||||
|
||||
fn error_response(&self) -> HttpResponse {
|
||||
let error_reason = match self {
|
||||
SessionInfoError::ServerError => "server error".to_string(),
|
||||
};
|
||||
|
||||
to_simple_json_error(&error_reason, self.status_code())
|
||||
}
|
||||
}
|
||||
|
||||
// NB: Not using derive_more::Display since Clion doesn't understand it.
|
||||
impl fmt::Display for SessionInfoError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{:?}", self)
|
||||
}
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
get,
|
||||
tag = "Users",
|
||||
path = "/v1/session",
|
||||
responses(
|
||||
(status = 200, description = "Get profile", body = SessionInfoSuccessResponse),
|
||||
(status = 500, description = "Server error", body = SessionInfoError),
|
||||
(status = 500, description = "Server error"),
|
||||
),
|
||||
)]
|
||||
pub async fn session_info_handler(
|
||||
@@ -64,100 +30,85 @@ pub async fn session_info_handler(
|
||||
mysql_pool: web::Data<MySqlPool>,
|
||||
session_checker: web::Data<SessionChecker>,
|
||||
avt_manager: web::Data<AvtCookieManager>,
|
||||
) -> Result<HttpResponse, SessionInfoError>
|
||||
) -> Result<HttpResponse, AdvancedCommonWebError>
|
||||
{
|
||||
let mut mysql_connection = mysql_pool.acquire()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("Could not acquire DB pool: {:?}", e);
|
||||
SessionInfoError::ServerError
|
||||
})?;
|
||||
let mut mysql_connection = mysql_pool.acquire().await?;
|
||||
|
||||
let maybe_user_session = session_checker
|
||||
.maybe_get_user_session_from_connection(&http_request, &mut mysql_connection)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("Session checker error: {:?}", e);
|
||||
SessionInfoError::ServerError
|
||||
.map_err(|err| {
|
||||
warn!("Error checking session: {:?}", err);
|
||||
err
|
||||
})?;
|
||||
|
||||
let mut logged_in = false;
|
||||
let mut user_info = None;
|
||||
|
||||
match maybe_user_session {
|
||||
None => {}
|
||||
Some(session_data) => {
|
||||
let feature_flags =
|
||||
UserSessionFeatureFlags::new(session_data.maybe_feature_flags.as_deref());
|
||||
if let Some(session_data) = maybe_user_session {
|
||||
let feature_flags =
|
||||
UserSessionFeatureFlags::new(session_data.maybe_feature_flags.as_deref());
|
||||
|
||||
if !session_data.is_banned {
|
||||
// NB: Banned users can't be logged in
|
||||
logged_in = true;
|
||||
user_info = Some(SessionUserInfo {
|
||||
core_info: UserDetailsLightBuilder::from_db_fields(
|
||||
&session_data.user_token,
|
||||
&session_data.username,
|
||||
&session_data.display_name,
|
||||
&session_data.email_gravatar_hash,
|
||||
),
|
||||
onboarding: SessionOnboardingState {
|
||||
email_not_set: session_data.email_is_synthetic,
|
||||
email_not_confirmed: !(session_data.email_confirmed || session_data.email_confirmed_by_google),
|
||||
password_not_set: session_data.is_without_password,
|
||||
username_not_customized: session_data.username_is_not_customized,
|
||||
},
|
||||
user_token: session_data.user_token,
|
||||
username: session_data.username.to_string(),
|
||||
display_name: session_data.display_name.to_string(),
|
||||
email_gravatar_hash: session_data.email_gravatar_hash.to_string(),
|
||||
if !session_data.is_banned {
|
||||
// NB: Banned users can't be logged in
|
||||
logged_in = true;
|
||||
user_info = Some(SessionUserInfo {
|
||||
core_info: UserDetailsLightBuilder::from_db_fields(
|
||||
&session_data.user_token,
|
||||
&session_data.username,
|
||||
&session_data.display_name,
|
||||
&session_data.email_gravatar_hash,
|
||||
),
|
||||
onboarding: SessionOnboardingState {
|
||||
email_not_set: session_data.email_is_synthetic,
|
||||
email_not_confirmed: !(session_data.email_confirmed || session_data.email_confirmed_by_google),
|
||||
password_not_set: session_data.is_without_password,
|
||||
username_not_customized: session_data.username_is_not_customized,
|
||||
},
|
||||
user_token: session_data.user_token,
|
||||
username: session_data.username.to_string(),
|
||||
display_name: session_data.display_name.to_string(),
|
||||
email_gravatar_hash: session_data.email_gravatar_hash.to_string(),
|
||||
|
||||
// Rollout / feature flags:
|
||||
can_access_studio: feature_flags.has_flag(UserFeatureFlag::Studio),
|
||||
maybe_feature_flags: feature_flags.clone_flags(),
|
||||
// Rollout / feature flags:
|
||||
can_access_studio: feature_flags.has_flag(UserFeatureFlag::Studio),
|
||||
maybe_feature_flags: feature_flags.clone_flags(),
|
||||
|
||||
// Premium plans:
|
||||
fakeyou_plan: FakeYouPlan::Free,
|
||||
storyteller_stream_plan: StorytellerStreamPlan::Free,
|
||||
// Premium plans:
|
||||
fakeyou_plan: FakeYouPlan::Free,
|
||||
storyteller_stream_plan: StorytellerStreamPlan::Free,
|
||||
|
||||
// Usage permissions:
|
||||
can_use_tts: session_data.can_use_tts,
|
||||
can_use_w2l: session_data.can_use_w2l,
|
||||
can_delete_own_tts_results: session_data.can_delete_own_tts_results,
|
||||
can_delete_own_w2l_results: session_data.can_delete_own_w2l_results,
|
||||
can_delete_own_account: session_data.can_delete_own_account,
|
||||
// Usage permissions:
|
||||
can_use_tts: session_data.can_use_tts,
|
||||
can_use_w2l: session_data.can_use_w2l,
|
||||
can_delete_own_tts_results: session_data.can_delete_own_tts_results,
|
||||
can_delete_own_w2l_results: session_data.can_delete_own_w2l_results,
|
||||
can_delete_own_account: session_data.can_delete_own_account,
|
||||
|
||||
// Contribution permissions:
|
||||
can_upload_tts_models: session_data.can_upload_tts_models,
|
||||
can_upload_w2l_templates: session_data.can_upload_w2l_templates,
|
||||
can_delete_own_tts_models: session_data.can_delete_own_tts_models,
|
||||
can_delete_own_w2l_templates: session_data.can_delete_own_w2l_templates,
|
||||
// Contribution permissions:
|
||||
can_upload_tts_models: session_data.can_upload_tts_models,
|
||||
can_upload_w2l_templates: session_data.can_upload_w2l_templates,
|
||||
can_delete_own_tts_models: session_data.can_delete_own_tts_models,
|
||||
can_delete_own_w2l_templates: session_data.can_delete_own_w2l_templates,
|
||||
|
||||
// Moderation permissions:
|
||||
can_approve_w2l_templates: session_data.can_approve_w2l_templates,
|
||||
can_edit_other_users_profiles: session_data.can_edit_other_users_profiles,
|
||||
can_edit_other_users_tts_models: session_data.can_edit_other_users_tts_models,
|
||||
can_edit_other_users_w2l_templates: session_data.can_edit_other_users_w2l_templates,
|
||||
can_delete_other_users_tts_models: session_data.can_delete_other_users_tts_models,
|
||||
can_delete_other_users_tts_results: session_data.can_delete_other_users_tts_results,
|
||||
can_delete_other_users_w2l_templates: session_data.can_delete_other_users_w2l_templates,
|
||||
can_delete_other_users_w2l_results: session_data.can_delete_other_users_w2l_results,
|
||||
can_ban_users: session_data.can_ban_users,
|
||||
can_delete_users: session_data.can_delete_users,
|
||||
});
|
||||
}
|
||||
// Moderation permissions:
|
||||
can_approve_w2l_templates: session_data.can_approve_w2l_templates,
|
||||
can_edit_other_users_profiles: session_data.can_edit_other_users_profiles,
|
||||
can_edit_other_users_tts_models: session_data.can_edit_other_users_tts_models,
|
||||
can_edit_other_users_w2l_templates: session_data.can_edit_other_users_w2l_templates,
|
||||
can_delete_other_users_tts_models: session_data.can_delete_other_users_tts_models,
|
||||
can_delete_other_users_tts_results: session_data.can_delete_other_users_tts_results,
|
||||
can_delete_other_users_w2l_templates: session_data.can_delete_other_users_w2l_templates,
|
||||
can_delete_other_users_w2l_results: session_data.can_delete_other_users_w2l_results,
|
||||
can_ban_users: session_data.can_ban_users,
|
||||
can_delete_users: session_data.can_delete_users,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let maybe_avt_cookie = match avt_manager.decode_cookie_payload_from_request(&http_request) {
|
||||
Ok(Some(_avt_cookie)) => None,
|
||||
_ => {
|
||||
let cookie = avt_manager.make_new_cookie()
|
||||
.map_err(|e| {
|
||||
warn!("avt cookie creation error: {:?}", e);
|
||||
SessionInfoError::ServerError
|
||||
})?;
|
||||
Some(cookie)
|
||||
}
|
||||
_ => Some(avt_manager.make_new_cookie()?),
|
||||
};
|
||||
|
||||
let response = SessionInfoSuccessResponse {
|
||||
@@ -166,8 +117,7 @@ pub async fn session_info_handler(
|
||||
user: user_info,
|
||||
};
|
||||
|
||||
let body = serde_json::to_string(&response)
|
||||
.map_err(|_e| SessionInfoError::ServerError)?;
|
||||
let body = serde_json::to_string(&response)?;
|
||||
|
||||
let mut response_builder = HttpResponse::Ok();
|
||||
|
||||
|
||||
@@ -1,2 +1,3 @@
|
||||
pub mod lookup;
|
||||
pub mod session_checker;
|
||||
pub mod session_checker;
|
||||
pub mod session_checker_error;
|
||||
@@ -5,7 +5,6 @@
|
||||
|
||||
use actix_artcraft::sessions::http_user_session_manager::HttpUserSessionManager;
|
||||
use actix_web::HttpRequest;
|
||||
use errors::AnyhowResult;
|
||||
use log::warn;
|
||||
use mysql_queries::queries::users::user_sessions::get_user_session_by_token::{get_user_session_by_token, get_user_session_by_token_pooled_connection, SessionUserRecord};
|
||||
use mysql_queries::queries::users::user_sessions::get_user_session_by_token_light::{get_user_session_by_token_light, SessionRecord};
|
||||
@@ -17,6 +16,7 @@ use sqlx::{Executor, MySql, MySqlPool};
|
||||
|
||||
use crate::http_server::session::lookup::user_session_extended::{UserSessionExtended, UserSessionPreferences, UserSessionPremiumPlanInfo, UserSessionRoleAndPermissions, UserSessionSubscriptionPlan, UserSessionUserDetails};
|
||||
use crate::http_server::session::lookup::user_session_feature_flags::UserSessionFeatureFlags;
|
||||
use crate::http_server::session::session_checker_error::SessionCheckerError;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct SessionChecker {
|
||||
@@ -40,7 +40,7 @@ impl SessionChecker {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_session_token(&self, request: &HttpRequest) -> AnyhowResult<Option<String>> {
|
||||
pub fn get_session_token(&self, request: &HttpRequest) -> Result<Option<String>, SessionCheckerError> {
|
||||
Ok(self.cookie_manager.decode_session_payload_from_request(request)?
|
||||
.map(|payload| payload.session_token))
|
||||
}
|
||||
@@ -56,7 +56,7 @@ impl SessionChecker {
|
||||
&self,
|
||||
request: &HttpRequest,
|
||||
pool: &MySqlPool
|
||||
) -> AnyhowResult<Option<SessionRecord>>
|
||||
) -> Result<Option<SessionRecord>, SessionCheckerError>
|
||||
{
|
||||
self.do_session_light_lookup_and_cookie_decode(request, pool).await
|
||||
}
|
||||
@@ -66,7 +66,7 @@ impl SessionChecker {
|
||||
&self,
|
||||
request: &HttpRequest,
|
||||
mysql_connection: &mut PoolConnection<MySql>,
|
||||
) -> AnyhowResult<Option<SessionRecord>>
|
||||
) -> Result<Option<SessionRecord>, SessionCheckerError>
|
||||
{
|
||||
self.do_session_light_lookup_and_cookie_decode(request, &mut **mysql_connection).await
|
||||
}
|
||||
@@ -76,7 +76,7 @@ impl SessionChecker {
|
||||
&self,
|
||||
request: &HttpRequest,
|
||||
mysql_executor: E,
|
||||
) -> AnyhowResult<Option<SessionRecord>>
|
||||
) -> Result<Option<SessionRecord>, SessionCheckerError>
|
||||
where E: 'e + Executor<'c, Database = MySql>
|
||||
{
|
||||
let maybe_session_token = self.cookie_manager.decode_session_payload_from_request(request)?
|
||||
@@ -95,18 +95,23 @@ impl SessionChecker {
|
||||
&self,
|
||||
mysql_executor: E,
|
||||
session_token: &str,
|
||||
) -> AnyhowResult<Option<SessionRecord>>
|
||||
) -> Result<Option<SessionRecord>, SessionCheckerError>
|
||||
where E: 'e + Executor<'c, Database = MySql>
|
||||
{
|
||||
match self.maybe_get_redis_cache_connection() {
|
||||
None => {
|
||||
get_user_session_by_token_light(mysql_executor, session_token).await
|
||||
Ok(get_user_session_by_token_light(mysql_executor, session_token).await?)
|
||||
}
|
||||
Some(mut redis_ttl_cache) => {
|
||||
let cache_key = RedisCacheKeys::session_record_light(session_token);
|
||||
redis_ttl_cache.lazy_load_if_not_cached(&cache_key, move || {
|
||||
// NB: Redis cache requires AnyhowResult; convert sqlx::Error to anyhow inside the closure,
|
||||
// then map any error from the cache layer back to SessionCheckerError::OtherError.
|
||||
let result = redis_ttl_cache.lazy_load_if_not_cached(&cache_key, move || async move {
|
||||
get_user_session_by_token_light(mysql_executor, session_token)
|
||||
}).await
|
||||
.await
|
||||
.map_err(anyhow::Error::from)
|
||||
}).await;
|
||||
result.map_err(SessionCheckerError::OtherError)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -119,7 +124,7 @@ impl SessionChecker {
|
||||
&self,
|
||||
request: &HttpRequest,
|
||||
pool: &MySqlPool,
|
||||
) -> AnyhowResult<Option<SessionUserRecord>>
|
||||
) -> Result<Option<SessionUserRecord>, SessionCheckerError>
|
||||
{
|
||||
self.do_user_session_lookup_and_cookie_decode(request, pool).await
|
||||
}
|
||||
@@ -129,7 +134,7 @@ impl SessionChecker {
|
||||
&self,
|
||||
request: &HttpRequest,
|
||||
mysql_connection: &mut PoolConnection<MySql>,
|
||||
) -> AnyhowResult<Option<SessionUserRecord>>
|
||||
) -> Result<Option<SessionUserRecord>, SessionCheckerError>
|
||||
{
|
||||
self.do_user_session_lookup_and_cookie_decode(request, &mut **mysql_connection).await
|
||||
}
|
||||
@@ -139,7 +144,7 @@ impl SessionChecker {
|
||||
&self,
|
||||
request: &HttpRequest,
|
||||
mysql_executor: E,
|
||||
) -> AnyhowResult<Option<SessionUserRecord>>
|
||||
) -> Result<Option<SessionUserRecord>, SessionCheckerError>
|
||||
where E: 'e + Executor<'c, Database = MySql>
|
||||
{
|
||||
let session_token = match self.get_session_token(request)? {
|
||||
@@ -155,18 +160,21 @@ impl SessionChecker {
|
||||
&self,
|
||||
mysql_executor: E,
|
||||
session_token: &str,
|
||||
) -> AnyhowResult<Option<SessionUserRecord>>
|
||||
) -> Result<Option<SessionUserRecord>, SessionCheckerError>
|
||||
where E: 'e + Executor<'c, Database = MySql>
|
||||
{
|
||||
match self.maybe_get_redis_cache_connection() {
|
||||
None => {
|
||||
get_user_session_by_token(mysql_executor, session_token).await
|
||||
Ok(get_user_session_by_token(mysql_executor, session_token).await?)
|
||||
}
|
||||
Some(mut redis_ttl_cache) => {
|
||||
let cache_key = RedisCacheKeys::session_record_user(session_token);
|
||||
redis_ttl_cache.lazy_load_if_not_cached(&cache_key, move || {
|
||||
let result = redis_ttl_cache.lazy_load_if_not_cached(&cache_key, move || async move {
|
||||
get_user_session_by_token(mysql_executor, session_token)
|
||||
}).await
|
||||
.await
|
||||
.map_err(anyhow::Error::from)
|
||||
}).await;
|
||||
result.map_err(SessionCheckerError::OtherError)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -179,7 +187,7 @@ impl SessionChecker {
|
||||
&self,
|
||||
request: &HttpRequest,
|
||||
pool: &MySqlPool,
|
||||
) -> AnyhowResult<Option<UserSessionExtended>>
|
||||
) -> Result<Option<UserSessionExtended>, SessionCheckerError>
|
||||
{
|
||||
let mut connection = pool.acquire().await?;
|
||||
self.maybe_get_user_session_extended_from_connection(request, &mut connection).await
|
||||
@@ -189,7 +197,7 @@ impl SessionChecker {
|
||||
&self,
|
||||
request: &HttpRequest,
|
||||
mysql_connection: &mut PoolConnection<MySql>,
|
||||
) -> AnyhowResult<Option<UserSessionExtended>>
|
||||
) -> Result<Option<UserSessionExtended>, SessionCheckerError>
|
||||
{
|
||||
let session_payload= match self.cookie_manager.decode_session_payload_from_request(request)? {
|
||||
None => return Ok(None),
|
||||
@@ -209,7 +217,8 @@ impl SessionChecker {
|
||||
list_active_user_subscriptions(
|
||||
mysql_connection,
|
||||
user_session.user_token.as_str()
|
||||
).await?;
|
||||
).await
|
||||
.map_err(SessionCheckerError::OtherError)?;
|
||||
|
||||
Ok(Some(UserSessionExtended {
|
||||
user_token: user_session.user_token.as_str().to_string(),
|
||||
@@ -283,4 +292,3 @@ impl SessionChecker {
|
||||
result.flatten()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
use std::error::Error;
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
use actix_artcraft::sessions::http_user_session_payload_error::HttpUserSessionPayloadError;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum SessionCheckerError {
|
||||
/// Error decoding or verifying the session.
|
||||
/// Rate of failures can point to the cause:
|
||||
/// - One-off errors: probably session forgeries
|
||||
/// - Lots of errors: maybe the wrong hmac secret
|
||||
SessionPayload(HttpUserSessionPayloadError),
|
||||
|
||||
/// Error looking up the session from the database.
|
||||
Sqlx(sqlx::Error),
|
||||
|
||||
/// Some other kind of error - perhaps with the Redis middleware.
|
||||
OtherError(anyhow::Error),
|
||||
}
|
||||
|
||||
impl Display for SessionCheckerError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::SessionPayload(e) => write!(f, "session payload error: {}", e),
|
||||
Self::Sqlx(e) => write!(f, "sqlx error: {}", e),
|
||||
Self::OtherError(e) => write!(f, "other error: {}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Error for SessionCheckerError {}
|
||||
|
||||
impl From<HttpUserSessionPayloadError> for SessionCheckerError {
|
||||
fn from(err: HttpUserSessionPayloadError) -> Self {
|
||||
Self::SessionPayload(err)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<sqlx::Error> for SessionCheckerError {
|
||||
fn from(err: sqlx::Error) -> Self {
|
||||
Self::Sqlx(err)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<anyhow::Error> for SessionCheckerError {
|
||||
fn from(err: anyhow::Error) -> Self {
|
||||
Self::OtherError(err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user