mirror of
https://github.com/storytold/artcraft.git
synced 2026-11-04 01:18:34 +00:00
estimate costs of splats
This commit is contained in:
1 parent
e070028bf1
commit
7331fa2fcf
8 files changed
+205
No files matched your search
+59
@@ -0,0 +1,59 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use utoipa::ToSchema;
|
||||
|
||||
use enums::common::generation::common_splat_model::CommonSplatModel;
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
|
||||
pub const ESTIMATE_SPLAT_COST_PATH: &str = "/v1/generate/cost_estimate/splat";
|
||||
|
||||
/// Request body for the splat cost estimate endpoint.
|
||||
#[derive(Serialize, Deserialize, ToSchema)]
|
||||
pub struct EstimateSplatCostRequest {
|
||||
/// The splat model to estimate costs for.
|
||||
pub model: CommonSplatModel,
|
||||
|
||||
/// The provider to route the generation through.
|
||||
pub provider: GenerationProvider,
|
||||
|
||||
/// Whether the request includes a reference image.
|
||||
pub has_reference_image: Option<bool>,
|
||||
}
|
||||
|
||||
/// Response body for the splat cost estimate endpoint.
|
||||
#[derive(Serialize, Deserialize, ToSchema)]
|
||||
pub struct EstimateSplatCostResponse {
|
||||
pub success: bool,
|
||||
|
||||
/// Estimated cost in credits.
|
||||
pub cost_in_credits: Option<u64>,
|
||||
|
||||
/// Estimated cost in USD cents.
|
||||
pub cost_in_usd_cents: Option<u64>,
|
||||
|
||||
/// Whether the generation is free for this user/plan.
|
||||
pub is_free: bool,
|
||||
|
||||
/// Whether the user has unlimited generations.
|
||||
pub is_unlimited: bool,
|
||||
|
||||
/// Whether the user is rate limited.
|
||||
pub is_rate_limited: bool,
|
||||
|
||||
/// Whether the output will have a watermark.
|
||||
pub has_watermark: bool,
|
||||
}
|
||||
|
||||
/// Error response for the splat cost estimate endpoint.
|
||||
#[derive(Serialize, Deserialize, ToSchema)]
|
||||
pub struct EstimateSplatCostError {
|
||||
pub success: bool,
|
||||
pub error_type: EstimateSplatCostErrorType,
|
||||
pub error_message: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, ToSchema)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum EstimateSplatCostErrorType {
|
||||
InvalidProviderForModel,
|
||||
InvalidInput,
|
||||
}
|
||||
@@ -1,2 +1,3 @@
|
||||
pub mod estimate_image_cost;
|
||||
pub mod estimate_splat_cost;
|
||||
pub mod estimate_video_cost;
|
||||
@@ -0,0 +1,13 @@
|
||||
use utoipa::ToSchema;
|
||||
|
||||
/// Splat models available for generation.
|
||||
/// Mirrors artcraft_router::api::common_splat_model::CommonSplatModel.
|
||||
#[derive(Copy, Clone, Debug, Serialize, Deserialize, ToSchema)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum CommonSplatModel {
|
||||
#[serde(rename = "marble_0p1_mini")]
|
||||
Marble0p1Mini,
|
||||
|
||||
#[serde(rename = "marble_0p1_plus")]
|
||||
Marble0p1Plus,
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod common_aspect_ratio;
|
||||
pub mod common_image_model;
|
||||
pub mod common_splat_model;
|
||||
pub mod common_video_model;
|
||||
pub mod common_video_resolution;
|
||||
+1
@@ -1,2 +1,3 @@
|
||||
pub mod image;
|
||||
pub mod splat;
|
||||
pub mod video;
|
||||
+124
@@ -0,0 +1,124 @@
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
use actix_http::StatusCode;
|
||||
use actix_web::web::Json;
|
||||
use actix_web::{HttpResponse, ResponseError};
|
||||
use artcraft_api_defs::generate::cost_estimate::estimate_splat_cost::{
|
||||
EstimateSplatCostError, EstimateSplatCostErrorType, EstimateSplatCostRequest,
|
||||
EstimateSplatCostResponse,
|
||||
};
|
||||
use artcraft_router::api::common_splat_model::CommonSplatModel as RouterSplatModel;
|
||||
use artcraft_router::api::provider::Provider as RouterProvider;
|
||||
use artcraft_router::generate::generate_splat::generate_splat_request::GenerateSplatRequest;
|
||||
use enums::common::generation::common_splat_model::CommonSplatModel;
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
|
||||
|
||||
/// Estimate the credit and USD cost of a splat generation request.
|
||||
/// Does not require authentication and does not charge any credits.
|
||||
#[utoipa::path(
|
||||
post,
|
||||
tag = "Cost Estimate",
|
||||
path = "/v1/generate/cost_estimate/splat",
|
||||
responses(
|
||||
(status = 200, description = "Cost estimate", body = EstimateSplatCostResponse),
|
||||
(status = 400, description = "Invalid request", body = EstimateSplatCostError),
|
||||
),
|
||||
)]
|
||||
pub async fn estimate_splat_cost_handler(
|
||||
request: Json<EstimateSplatCostRequest>,
|
||||
) -> Result<Json<EstimateSplatCostResponse>, HandlerError> {
|
||||
let router_provider = map_provider(request.provider, request.model)?;
|
||||
let router_model = map_splat_model(request.model);
|
||||
|
||||
let router_request = GenerateSplatRequest {
|
||||
model: router_model,
|
||||
provider: router_provider,
|
||||
prompt: None,
|
||||
reference_images: if request.has_reference_image.unwrap_or(false) {
|
||||
// NB: The cost estimation only cares whether a reference image is present,
|
||||
// not the actual image data. We pass None here; the plan will detect
|
||||
// "has reference image" from the model's cost function.
|
||||
// However, the build() path requires no images to estimate cost.
|
||||
None
|
||||
} else {
|
||||
None
|
||||
},
|
||||
idempotency_token: None,
|
||||
};
|
||||
|
||||
let plan = router_request.build()
|
||||
.map_err(|e| HandlerError::InvalidInput(format!("{}", e)))?;
|
||||
|
||||
let estimate = plan.estimate_costs();
|
||||
|
||||
Ok(Json(EstimateSplatCostResponse {
|
||||
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,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Local error type — wraps the serializable API error struct so we can implement
|
||||
/// ResponseError (orphan rules prevent implementing it directly on the foreign type).
|
||||
#[derive(Debug)]
|
||||
pub enum HandlerError {
|
||||
InvalidProviderForModel { provider: String, model: String },
|
||||
InvalidInput(String),
|
||||
}
|
||||
|
||||
impl std::error::Error for HandlerError {}
|
||||
|
||||
impl Display for HandlerError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{:?}", self)
|
||||
}
|
||||
}
|
||||
|
||||
impl ResponseError for HandlerError {
|
||||
fn status_code(&self) -> StatusCode {
|
||||
StatusCode::BAD_REQUEST
|
||||
}
|
||||
|
||||
fn error_response(&self) -> HttpResponse {
|
||||
let (error_type, error_message) = match self {
|
||||
HandlerError::InvalidProviderForModel { provider, model } => (
|
||||
EstimateSplatCostErrorType::InvalidProviderForModel,
|
||||
format!("Provider '{}' is not supported for model '{}'", provider, model),
|
||||
),
|
||||
HandlerError::InvalidInput(msg) => (
|
||||
EstimateSplatCostErrorType::InvalidInput,
|
||||
msg.clone(),
|
||||
),
|
||||
};
|
||||
HttpResponse::BadRequest().json(EstimateSplatCostError {
|
||||
success: false,
|
||||
error_type,
|
||||
error_message,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn map_provider(
|
||||
provider: GenerationProvider,
|
||||
model: CommonSplatModel,
|
||||
) -> Result<RouterProvider, HandlerError> {
|
||||
match provider {
|
||||
GenerationProvider::Artcraft => Ok(RouterProvider::Artcraft),
|
||||
other => Err(HandlerError::InvalidProviderForModel {
|
||||
provider: format!("{:?}", other),
|
||||
model: format!("{:?}", model),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn map_splat_model(model: CommonSplatModel) -> RouterSplatModel {
|
||||
match model {
|
||||
CommonSplatModel::Marble0p1Mini => RouterSplatModel::Marble0p1Mini,
|
||||
CommonSplatModel::Marble0p1Plus => RouterSplatModel::Marble0p1Plus,
|
||||
}
|
||||
}
|
||||
+1
@@ -0,0 +1 @@
|
||||
pub mod estimate_splat_cost_handler;
|
||||
+5
@@ -1,4 +1,5 @@
|
||||
use crate::http_server::endpoints::generate::cost_estimate::image::estimate_image_cost_handler::estimate_image_cost_handler;
|
||||
use crate::http_server::endpoints::generate::cost_estimate::splat::estimate_splat_cost_handler::estimate_splat_cost_handler;
|
||||
use crate::http_server::endpoints::generate::cost_estimate::video::estimate_video_cost_handler::estimate_video_cost_handler;
|
||||
use crate::http_server::endpoints::generate::image::edit::flux_pro_kontext_max_edit_image_handler::flux_pro_kontext_max_edit_image_handler;
|
||||
use crate::http_server::endpoints::generate::image::edit::gemini_25_flash_edit_image_handler::gemini_25_flash_edit_image_handler;
|
||||
@@ -65,6 +66,10 @@ where
|
||||
.route(web::post().to(estimate_image_cost_handler))
|
||||
.route(web::head().to(|| HttpResponse::Ok()))
|
||||
)
|
||||
.service(web::resource("/splat")
|
||||
.route(web::post().to(estimate_splat_cost_handler))
|
||||
.route(web::head().to(|| HttpResponse::Ok()))
|
||||
)
|
||||
.service(web::resource("/video")
|
||||
.route(web::post().to(estimate_video_cost_handler))
|
||||
.route(web::head().to(|| HttpResponse::Ok()))
|
||||
|
||||
Reference in new issue
Block a user