mirror of
https://github.com/storytold/artcraft.git
synced 2026-10-21 01:12:40 -04:00
Compare commits
23
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3dbc540fb2 | ||
|
|
f3943818d9 | ||
|
|
0a99b77a6a | ||
|
|
9cd6685f02 | ||
|
|
89e8346fa6 | ||
|
|
79a85e1b32 | ||
|
|
fe1235a93c | ||
|
|
5217e9ff07 | ||
|
|
bb375e405f | ||
|
|
0a44642c95 | ||
|
|
0b68e04b5f | ||
|
|
a0e6b5721b | ||
|
|
8ffa1de27f | ||
|
|
1f0eb35ff1 | ||
|
|
43ec1ae5d7 | ||
|
|
625ba41423 | ||
|
|
2a62bcc7f8 | ||
|
|
5c92adff8f | ||
|
|
1276df9f29 | ||
|
|
026d356d9c | ||
|
|
2bfeb49d17 | ||
|
|
4752c54851 | ||
|
|
e0779897b6 |
@@ -9,6 +9,7 @@ on:
|
||||
# NB: Default-branch doesn't work, despite Github's documentation.
|
||||
#- $default-branch
|
||||
- main
|
||||
- test
|
||||
env:
|
||||
IMAGE_NAME: storyteller-rust
|
||||
|
||||
|
||||
@@ -35,6 +35,9 @@ script/socat_dev_database_connect.sh
|
||||
# Python virtual envs (named by convention)
|
||||
venv/
|
||||
|
||||
# Testing
|
||||
/requests/
|
||||
|
||||
# Scratch data
|
||||
scratch/
|
||||
test_data/temp/
|
||||
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
{
|
||||
"db_name": "MySQL",
|
||||
"query": "\nUPDATE generic_inference_jobs\nSET\n status = ?,\n failure_reason = ?,\n internal_debugging_failure_reason = ?,\n frontend_failure_category = ?,\n retry_at = NULL\nWHERE token = ?\n ",
|
||||
"describe": {
|
||||
"columns": [],
|
||||
"parameters": {
|
||||
"Right": 5
|
||||
},
|
||||
"nullable": []
|
||||
},
|
||||
"hash": "15939605e0da2722693371799ba94c7d0cc2f77a5c5b37e346855996d16033ec"
|
||||
}
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
{
|
||||
"db_name": "MySQL",
|
||||
"query": "\nINSERT INTO generic_inference_jobs\nSET\n token = ?,\n uuid_idempotency_token = ?,\n\n job_type = ?,\n\n maybe_external_third_party = ?,\n maybe_external_third_party_id = ?,\n\n product_category = ?,\n inference_category = ?,\n\n maybe_model_type = NULL,\n maybe_model_token = NULL,\n\n maybe_input_source_token = NULL,\n maybe_input_source_token_type = NULL,\n\n maybe_download_url = NULL,\n maybe_cover_image_media_file_token = NULL,\n\n maybe_prompt_token = ?,\n\n maybe_raw_inference_text = NULL,\n\n maybe_inference_args = ?,\n\n maybe_creator_user_token = ?,\n maybe_creator_anonymous_visitor_token = ?,\n creator_ip_address = ?,\n creator_set_visibility = ?,\n\n priority_level = 0,\n is_keepalive_required = FALSE,\n max_duration_seconds = 0,\n\n is_debug_request = FALSE,\n maybe_routing_tag = NULL,\n\n status = ?\n ",
|
||||
"describe": {
|
||||
"columns": [],
|
||||
"parameters": {
|
||||
"Right": 14
|
||||
},
|
||||
"nullable": []
|
||||
},
|
||||
"hash": "33aa8a1cf85bf1b8b68f4e2ab46541f25d7e8bb5eeb503c5ec14cb5e115f69ce"
|
||||
}
|
||||
+91
@@ -0,0 +1,91 @@
|
||||
{
|
||||
"db_name": "MySQL",
|
||||
"query": "\nSELECT\n jobs.token as `job_token: tokens::tokens::generic_inference_jobs::InferenceJobToken`,\n jobs.maybe_external_third_party_id as `order_id`,\n jobs.maybe_creator_user_token as `maybe_creator_user_token: tokens::tokens::users::UserToken`,\n jobs.maybe_creator_anonymous_visitor_token as `maybe_creator_anonymous_visitor_token: tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken`,\n jobs.creator_ip_address,\n jobs.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,\n jobs.maybe_prompt_token as `maybe_prompt_token: tokens::tokens::prompts::PromptToken`\n\nFROM generic_inference_jobs as jobs\n\nWHERE jobs.maybe_external_third_party = ?\n AND jobs.status NOT IN ('complete_success', 'complete_failure')\n AND jobs.maybe_external_third_party_id IS NOT NULL\n\nLIMIT 25000\n ",
|
||||
"describe": {
|
||||
"columns": [
|
||||
{
|
||||
"ordinal": 0,
|
||||
"name": "job_token: tokens::tokens::generic_inference_jobs::InferenceJobToken",
|
||||
"type_info": {
|
||||
"type": "VarString",
|
||||
"flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE",
|
||||
"char_set": 224,
|
||||
"max_size": 128
|
||||
}
|
||||
},
|
||||
{
|
||||
"ordinal": 1,
|
||||
"name": "order_id",
|
||||
"type_info": {
|
||||
"type": "VarString",
|
||||
"flags": "MULTIPLE_KEY | BINARY",
|
||||
"char_set": 224,
|
||||
"max_size": 256
|
||||
}
|
||||
},
|
||||
{
|
||||
"ordinal": 2,
|
||||
"name": "maybe_creator_user_token: tokens::tokens::users::UserToken",
|
||||
"type_info": {
|
||||
"type": "VarString",
|
||||
"flags": "MULTIPLE_KEY | BINARY",
|
||||
"char_set": 224,
|
||||
"max_size": 128
|
||||
}
|
||||
},
|
||||
{
|
||||
"ordinal": 3,
|
||||
"name": "maybe_creator_anonymous_visitor_token: tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken",
|
||||
"type_info": {
|
||||
"type": "VarString",
|
||||
"flags": "MULTIPLE_KEY | BINARY",
|
||||
"char_set": 224,
|
||||
"max_size": 128
|
||||
}
|
||||
},
|
||||
{
|
||||
"ordinal": 4,
|
||||
"name": "creator_ip_address",
|
||||
"type_info": {
|
||||
"type": "VarString",
|
||||
"flags": "NOT_NULL | MULTIPLE_KEY | BINARY | NO_DEFAULT_VALUE",
|
||||
"char_set": 224,
|
||||
"max_size": 160
|
||||
}
|
||||
},
|
||||
{
|
||||
"ordinal": 5,
|
||||
"name": "creator_set_visibility: enums::common::visibility::Visibility",
|
||||
"type_info": {
|
||||
"type": "String",
|
||||
"flags": "NOT_NULL | BINARY | ENUM",
|
||||
"char_set": 224,
|
||||
"max_size": 28
|
||||
}
|
||||
},
|
||||
{
|
||||
"ordinal": 6,
|
||||
"name": "maybe_prompt_token: tokens::tokens::prompts::PromptToken",
|
||||
"type_info": {
|
||||
"type": "VarString",
|
||||
"flags": "BINARY",
|
||||
"char_set": 224,
|
||||
"max_size": 128
|
||||
}
|
||||
}
|
||||
],
|
||||
"parameters": {
|
||||
"Right": 1
|
||||
},
|
||||
"nullable": [
|
||||
false,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
]
|
||||
},
|
||||
"hash": "ca96942d8214edb84f163cf40131ceed873209e86718054e1502219fc81d2dc6"
|
||||
}
|
||||
Generated
+425
-317
File diff suppressed because it is too large
Load Diff
+12
-8
@@ -16,6 +16,7 @@ members = [
|
||||
"crates/api_clients/midjourney_client",
|
||||
"crates/api_clients/openai_sora_client",
|
||||
"crates/api_clients/openart_client",
|
||||
"crates/api_clients/seedance2pro",
|
||||
"crates/api_clients/runwayml_client",
|
||||
"crates/api_clients/thumbnail_generator",
|
||||
"crates/api_clients/world_labs_client",
|
||||
@@ -116,6 +117,7 @@ members = [
|
||||
"crates/service/job/email_sender_job",
|
||||
"crates/service/job/es-update-job",
|
||||
"crates/service/job/inference_job",
|
||||
"crates/service/job/seedance2_pro_job",
|
||||
"crates/service/job/tts_download_job",
|
||||
"crates/service/job/tts_inference_job",
|
||||
"crates/service/job/w2l_download_job",
|
||||
@@ -232,6 +234,7 @@ memory_store = { path = "crates/lib/caching/memory_store" }
|
||||
midjourney_client = { path = "crates/api_clients/midjourney_client" }
|
||||
mimetypes = { path = "crates/lib/files/mimetypes" }
|
||||
openai_sora_client = { path = "crates/api_clients/openai_sora_client" }
|
||||
seedance2pro = { path = "crates/api_clients/seedance2pro" }
|
||||
primitives = { path = "crates/lib/data/primitives" }
|
||||
shared_service_components = { path = "crates/lib/shared_service_components" }
|
||||
sqlite_tasks = { path = "crates/schema/database/sqlite_tasks" }
|
||||
@@ -258,24 +261,24 @@ actix-utils = "3.0.1"
|
||||
actix-web = "4.7.0" # TODO(bt): We're behind and need to bump versions.
|
||||
actix-web-actors = "4.3.1+deprecated"
|
||||
actix-web-lab = "0.20.2"
|
||||
anyhow = "1.0.98"
|
||||
anyhow = "1.0.102"
|
||||
base64 = "0.22.1"
|
||||
bstr = "1.12.1"
|
||||
bytes = "1.10.1"
|
||||
chrono = { version = "0.4.41", features = ["serde"] }
|
||||
cookie = "0.18.1"
|
||||
env = "1.0.1"
|
||||
env_logger = "0.11.8"
|
||||
env_logger = "0.11.9"
|
||||
futures = { version = "0.3.31" }
|
||||
hex-literal = "1.0.0"
|
||||
indexmap = "2.9.0"
|
||||
log = "0.4.27"
|
||||
log = "0.4.29"
|
||||
md-5 = "0.10.6" # insecure but used for gravatar
|
||||
once_cell = "1.21.3"
|
||||
r2d2 = "0.8.10"
|
||||
rand = "0.9.1"
|
||||
redis = { version = "1.0.2", features = ["r2d2", "tokio-rustls-comp", "tokio-comp", "cluster-async", "aio"] }
|
||||
regex = "1.12.2"
|
||||
regex = "1.12.3"
|
||||
reqwest = { version = "0.12.23", default-features = false, features = [ "charset", "http2", "system-proxy", "json", "cookies", "rustls-tls", "gzip" ] }
|
||||
reqwest-websocket = { version = "0.5.0" }
|
||||
serde = { version = "1.0.219", features = ["derive"] }
|
||||
@@ -293,11 +296,11 @@ tempdir = "0.3.7"
|
||||
tempfile = "3.20.0"
|
||||
thiserror = "2.0.12"
|
||||
tokio = { version = "1.45.1", features = ["macros"] }
|
||||
toml = "0.8.22"
|
||||
toml = "0.8.23"
|
||||
url = { version = "2.5.4", features = ["serde"] }
|
||||
uuid = "1.18.1"
|
||||
|
||||
# NB: wreq gets past fingerprinting, but we have to deal with platform-specific build issues.
|
||||
# NB(1): wreq gets past fingerprinting, but we have to deal with platform-specific build issues.
|
||||
#
|
||||
# Mac:
|
||||
# brew install --formula cmake pkg-config llvm
|
||||
@@ -307,5 +310,6 @@ uuid = "1.18.1"
|
||||
# choco install pkgconfiglite -y
|
||||
# choco install llvm -y
|
||||
# choco install nasm -y
|
||||
wreq = { version = "6.0.0-rc.20", features = ["json", "multipart", "cookies", "ws", "stream"] }
|
||||
wreq-util = "3.0.0-rc.6"
|
||||
wreq-util = "3.0.0-rc.10"
|
||||
wreq = { version = "6.0.0-rc.28", features = ["json", "query", "multipart", "cookies", "ws", "stream", "gzip", "form"] }
|
||||
|
||||
|
||||
+1
-1
@@ -3,7 +3,7 @@
|
||||
-- noinspection SqlResolveForFile
|
||||
|
||||
ALTER TABLE media_files
|
||||
ADD COLUMN origin_product_category VARCHAR(16) NOT NULL DEFAULT "unknown"
|
||||
ADD COLUMN origin_product_category VARCHAR(16) NOT NULL DEFAULT 'unknown'
|
||||
AFTER origin_category;
|
||||
|
||||
ALTER TABLE media_files
|
||||
|
||||
+1
-1
@@ -3,7 +3,7 @@
|
||||
-- noinspection SqlResolveForFile
|
||||
|
||||
ALTER TABLE media_files
|
||||
ADD COLUMN media_class VARCHAR(16) NOT NULL DEFAULT "unknown"
|
||||
ADD COLUMN media_class VARCHAR(16) NOT NULL DEFAULT 'unknown'
|
||||
AFTER media_type;
|
||||
|
||||
ALTER TABLE media_files
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
-- noinspection SqlResolveForFile
|
||||
|
||||
ALTER TABLE media_files
|
||||
ADD COLUMN nsfw_status VARCHAR(32) NOT NULL DEFAULT "unknown"
|
||||
ADD COLUMN nsfw_status VARCHAR(32) NOT NULL DEFAULT 'unknown'
|
||||
AFTER maybe_scene_source_media_file_token;
|
||||
|
||||
ALTER TABLE media_files
|
||||
|
||||
@@ -33,7 +33,7 @@ RUN apt-get update \
|
||||
FROM ubuntu:jammy as rust-base
|
||||
|
||||
# NB: This can be "stable" or another version.
|
||||
ARG RUST_TOOLCHAIN="1.86.0"
|
||||
ARG RUST_TOOLCHAIN="1.93.0"
|
||||
|
||||
WORKDIR /tmp
|
||||
|
||||
@@ -48,8 +48,12 @@ RUN apt-get update \
|
||||
curl \
|
||||
ffmpeg \
|
||||
fontconfig \
|
||||
git \
|
||||
libclang-dev \
|
||||
libfontconfig1-dev \
|
||||
libssl-dev \
|
||||
musl-tools \
|
||||
perl \
|
||||
pkg-config
|
||||
|
||||
# NB: Fix for fontconfig (servo-fontconfig-sys): https://github.com/alacritty/alacritty/issues/4423#issuecomment-727277235
|
||||
@@ -133,6 +137,12 @@ RUN SQLX_OFFLINE=true \
|
||||
--release \
|
||||
--bin storyteller-web
|
||||
|
||||
RUN SQLX_OFFLINE=true \
|
||||
LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH} \
|
||||
$HOME/.cargo/bin/cargo build \
|
||||
--release \
|
||||
--bin seedance2-pro-job
|
||||
|
||||
RUN SQLX_OFFLINE=true \
|
||||
LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH} \
|
||||
$HOME/.cargo/bin/cargo build \
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
pub mod kling_2_5_turbo_multi_function_video_gen;
|
||||
pub mod kling_2_6_multi_function_video_gen;
|
||||
pub mod seedance_2p0_multi_function_video_gen;
|
||||
pub mod sora_2_multi_function_video_gen;
|
||||
pub mod sora_2_pro_multi_function_video_gen;
|
||||
pub mod veo_3p1_fast_multi_function_video_gen;
|
||||
|
||||
+68
@@ -0,0 +1,68 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokens::tokens::generic_inference_jobs::InferenceJobToken;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
use utoipa::ToSchema;
|
||||
|
||||
pub const SEEDANCE_2P0_MULTI_FUNCTION_VIDEO_GEN_PATH: &str = "/v1/generate/video/multi_function/seedance_2p0";
|
||||
|
||||
/// Seedance 2.0 multi-function video generation (text-to-video, image-to-video, and reference-based).
|
||||
#[derive(Serialize, Deserialize, ToSchema)]
|
||||
pub struct Seedance2p0MultiFunctionVideoGenRequest {
|
||||
/// Idempotency token to prevent duplicate requests.
|
||||
pub uuid_idempotency_token: String,
|
||||
|
||||
/// Text prompt.
|
||||
pub prompt: Option<String>,
|
||||
|
||||
/// Optional start frame image.
|
||||
pub start_frame_media_token: Option<MediaFileToken>,
|
||||
|
||||
/// Optional end frame image.
|
||||
pub end_frame_media_token: Option<MediaFileToken>,
|
||||
|
||||
/// Optional reference images (reference mode). When present, takes priority over start/end frames.
|
||||
pub reference_image_media_tokens: Option<Vec<MediaFileToken>>,
|
||||
|
||||
/// Video aspect ratio / resolution.
|
||||
pub aspect_ratio: Option<Seedance2p0AspectRatio>,
|
||||
|
||||
/// Duration in seconds (4–15).
|
||||
pub duration_seconds: Option<u8>,
|
||||
|
||||
/// Number of videos to generate.
|
||||
pub batch_count: Option<Seedance2p0BatchCount>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, ToSchema)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum Seedance2p0AspectRatio {
|
||||
/// 16:9 landscape (1280x720)
|
||||
Landscape16x9,
|
||||
/// 9:16 portrait (720x1280)
|
||||
Portrait9x16,
|
||||
/// 1:1 square (720x720)
|
||||
Square1x1,
|
||||
/// 4:3 standard (960x720)
|
||||
Standard4x3,
|
||||
/// 3:4 portrait (720x960)
|
||||
Portrait3x4,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, ToSchema)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum Seedance2p0BatchCount {
|
||||
One,
|
||||
Two,
|
||||
Four,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, ToSchema)]
|
||||
pub struct Seedance2p0MultiFunctionVideoGenResponse {
|
||||
pub success: bool,
|
||||
|
||||
/// The first inference job token (always present).
|
||||
pub inference_job_token: InferenceJobToken,
|
||||
|
||||
/// All inference job tokens for the batch (includes `inference_job_token` as the first entry).
|
||||
pub all_inference_job_tokens: Vec<InferenceJobToken>,
|
||||
}
|
||||
+1
@@ -1,4 +1,5 @@
|
||||
pub mod kling_2p5_turbo_pro_multi_function_video_gen;
|
||||
pub mod seedance_2p0_multi_function_video_gen;
|
||||
pub mod kling_2p6_pro_multi_function_video_gen;
|
||||
pub mod sora_2_multi_function_video_gen;
|
||||
pub mod sora_2_pro_multi_function_video_gen;
|
||||
|
||||
+18
@@ -0,0 +1,18 @@
|
||||
use crate::credentials::storyteller_credential_set::StorytellerCredentialSet;
|
||||
use crate::error::storyteller_error::StorytellerError;
|
||||
use crate::utils::api_host::ApiHost;
|
||||
use crate::utils::basic_json_post_request::basic_json_post_request;
|
||||
use artcraft_api_defs::generate::video::multi_function::seedance_2p0_multi_function_video_gen::{Seedance2p0MultiFunctionVideoGenRequest, Seedance2p0MultiFunctionVideoGenResponse, SEEDANCE_2P0_MULTI_FUNCTION_VIDEO_GEN_PATH};
|
||||
|
||||
pub async fn seedance_2p0_multi_function_video_gen(
|
||||
api_host: &ApiHost,
|
||||
maybe_creds: Option<&StorytellerCredentialSet>,
|
||||
request: Seedance2p0MultiFunctionVideoGenRequest,
|
||||
) -> Result<Seedance2p0MultiFunctionVideoGenResponse, StorytellerError> {
|
||||
Ok(basic_json_post_request(
|
||||
api_host,
|
||||
SEEDANCE_2P0_MULTI_FUNCTION_VIDEO_GEN_PATH,
|
||||
maybe_creds,
|
||||
request,
|
||||
).await?)
|
||||
}
|
||||
@@ -36,11 +36,11 @@ sha3.workspace = true
|
||||
thiserror.workspace = true
|
||||
#tokio.workspace = true
|
||||
toml.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
wreq-util.workspace = true
|
||||
wreq.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
testing = { path = "../../lib/testing" }
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
[package]
|
||||
name = "seedance2pro"
|
||||
edition = "2024"
|
||||
version = "0.0.1"
|
||||
publish = false
|
||||
|
||||
[package.metadata.cargo-udeps.ignore]
|
||||
normal = ["workspace-build-acceleration"]
|
||||
|
||||
[lib]
|
||||
name = "seedance2pro"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[dependencies]
|
||||
# Workspace hack (faster compile times)
|
||||
workspace-build-acceleration.workspace = true
|
||||
|
||||
# Internal
|
||||
errors.workspace = true
|
||||
|
||||
# External
|
||||
chrono.workspace = true
|
||||
cloudflare_errors.workspace = true
|
||||
log.workspace = true
|
||||
rand.workspace = true
|
||||
serde.workspace = true
|
||||
serde_derive.workspace = true
|
||||
serde_json.workspace = true
|
||||
wreq-util.workspace = true
|
||||
wreq.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
chrono.workspace = true
|
||||
env.workspace = true
|
||||
env_logger.workspace = true
|
||||
testing = { path = "../../lib/testing" }
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod seedance2pro_cookies;
|
||||
pub mod seedance2pro_session;
|
||||
@@ -0,0 +1,24 @@
|
||||
/// Cookies are the credential required to interact with the Seedance2 Pro API.
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Seedance2ProCookies {
|
||||
cookies: String,
|
||||
}
|
||||
|
||||
impl Seedance2ProCookies {
|
||||
pub fn new(cookies: String) -> Self {
|
||||
Seedance2ProCookies { cookies }
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.cookies
|
||||
}
|
||||
|
||||
pub fn as_bytes(&self) -> &[u8] {
|
||||
self.cookies.as_bytes()
|
||||
}
|
||||
|
||||
pub fn to_string(&self) -> String {
|
||||
self.cookies.clone()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
use crate::creds::seedance2pro_cookies::Seedance2ProCookies;
|
||||
|
||||
/// Holds the full session data needed to make authenticated requests to Seedance2 Pro.
|
||||
#[derive(Clone)]
|
||||
pub struct Seedance2ProSession {
|
||||
pub cookies: Seedance2ProCookies,
|
||||
}
|
||||
|
||||
impl Seedance2ProSession {
|
||||
pub fn new(cookies: Seedance2ProCookies) -> Self {
|
||||
Seedance2ProSession { cookies }
|
||||
}
|
||||
|
||||
pub fn from_cookies_string(cookies: String) -> Self {
|
||||
Seedance2ProSession {
|
||||
cookies: Seedance2ProCookies::new(cookies),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
pub mod seedance2pro_client_error;
|
||||
pub mod seedance2pro_error;
|
||||
pub mod seedance2pro_generic_api_error;
|
||||
pub mod seedance2pro_specific_api_error;
|
||||
@@ -0,0 +1,26 @@
|
||||
use std::error::Error;
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Seedance2ProClientError {
|
||||
/// No cookies are present in the session.
|
||||
NoCookiesPresent,
|
||||
|
||||
/// Error parsing a URL.
|
||||
UrlParseError(url::ParseError),
|
||||
|
||||
/// An error was encountered in building the Wreq client.
|
||||
WreqClientError(wreq::Error),
|
||||
}
|
||||
|
||||
impl Error for Seedance2ProClientError {}
|
||||
|
||||
impl Display for Seedance2ProClientError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::NoCookiesPresent => write!(f, "No cookies present in the session."),
|
||||
Self::UrlParseError(err) => write!(f, "URL parse error: {}", err),
|
||||
Self::WreqClientError(err) => write!(f, "Wreq client error (during client creation): {}", err),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
use crate::error::seedance2pro_client_error::Seedance2ProClientError;
|
||||
use crate::error::seedance2pro_generic_api_error::Seedance2ProGenericApiError;
|
||||
use crate::error::seedance2pro_specific_api_error::Seedance2ProSpecificApiError;
|
||||
use cloudflare_errors::cloudflare_error::CloudflareError;
|
||||
use std::error::Error;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Seedance2ProError {
|
||||
Client(Seedance2ProClientError),
|
||||
ApiSpecific(Seedance2ProSpecificApiError),
|
||||
ApiGeneric(Seedance2ProGenericApiError),
|
||||
}
|
||||
|
||||
impl Seedance2ProError {
|
||||
pub fn is_having_downtime_issues(&self) -> bool {
|
||||
match self {
|
||||
Self::ApiGeneric(Seedance2ProGenericApiError::Http502ErrorBadGateway(_)) => true,
|
||||
Self::ApiGeneric(Seedance2ProGenericApiError::CloudflareError(CloudflareError::BadGateway502)) => true,
|
||||
Self::ApiGeneric(Seedance2ProGenericApiError::CloudflareError(CloudflareError::GatewayTimeout504)) => true,
|
||||
Self::ApiGeneric(Seedance2ProGenericApiError::CloudflareError(CloudflareError::TimeoutOccurred524)) => true,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Error for Seedance2ProError {}
|
||||
|
||||
impl std::fmt::Display for Seedance2ProError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Client(e) => write!(f, "Seedance2ProClientError: {:?}", e),
|
||||
Self::ApiSpecific(e) => write!(f, "Seedance2ProSpecificApiError: {:?}", e),
|
||||
Self::ApiGeneric(e) => write!(f, "Seedance2ProGenericApiError: {:?}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Seedance2ProClientError> for Seedance2ProError {
|
||||
fn from(error: Seedance2ProClientError) -> Self {
|
||||
Self::Client(error)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Seedance2ProSpecificApiError> for Seedance2ProError {
|
||||
fn from(error: Seedance2ProSpecificApiError) -> Self {
|
||||
Self::ApiSpecific(error)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Seedance2ProGenericApiError> for Seedance2ProError {
|
||||
fn from(error: Seedance2ProGenericApiError) -> Self {
|
||||
Self::ApiGeneric(error)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
use cloudflare_errors::cloudflare_error::CloudflareError;
|
||||
use std::error::Error;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use wreq::StatusCode;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Seedance2ProGenericApiError {
|
||||
/// Specific Cloudflare errors.
|
||||
CloudflareError(CloudflareError),
|
||||
|
||||
/// HTTP 502 Bad Gateway.
|
||||
/// We preserve the message for debugging.
|
||||
Http502ErrorBadGateway(String),
|
||||
|
||||
/// serde_json::Error, likely from JSON deserialization schema mismatch.
|
||||
/// Includes the original body.
|
||||
SerdeResponseParseErrorWithBody(serde_json::Error, String),
|
||||
|
||||
/// serde_json::Error on a non-200 (error) response.
|
||||
/// Includes the original body.
|
||||
SerdeParseErrorWithBodyOnNon200(serde_json::Error, String),
|
||||
|
||||
/// An uncategorized bad HTTP response.
|
||||
UncategorizedBadResponse(String),
|
||||
|
||||
/// An uncategorized bad HTTP response with status code and body.
|
||||
UncategorizedBadResponseWithStatusAndBody {
|
||||
status_code: StatusCode,
|
||||
body: String,
|
||||
},
|
||||
|
||||
/// An uncaught error from the HTTP client.
|
||||
WreqError(wreq::Error),
|
||||
}
|
||||
|
||||
impl Error for Seedance2ProGenericApiError {}
|
||||
|
||||
impl Display for Seedance2ProGenericApiError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::CloudflareError(err) => write!(f, "Cloudflare error: {}", err),
|
||||
Self::Http502ErrorBadGateway(msg) => write!(f, "HTTP 502 Bad Gateway: {}", msg),
|
||||
Self::SerdeResponseParseErrorWithBody(err, body) => write!(f, "Failed to parse response body: {:?}. Body: {}", err, body),
|
||||
Self::SerdeParseErrorWithBodyOnNon200(err, body) => write!(f, "Failed to parse non-200 response body: {:?}. Body: {}", err, body),
|
||||
Self::UncategorizedBadResponse(msg) => write!(f, "Uncategorized bad response: {}", msg),
|
||||
Self::UncategorizedBadResponseWithStatusAndBody { status_code, body } => write!(f, "Uncategorized bad response: status code {}, body: {}", status_code, body),
|
||||
Self::WreqError(err) => write!(f, "Wreq client error: {}", err),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
use std::error::Error;
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Seedance2ProSpecificApiError {
|
||||
/// The session cookies are expired or invalid.
|
||||
UnauthorizedSessionExpired,
|
||||
|
||||
/// The video generation request was flagged as a content violation.
|
||||
VideoGenerationViolation(String),
|
||||
}
|
||||
|
||||
impl Error for Seedance2ProSpecificApiError {}
|
||||
|
||||
impl Display for Seedance2ProSpecificApiError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::UnauthorizedSessionExpired => write!(f, "Unauthorized: session cookies expired or invalid."),
|
||||
Self::VideoGenerationViolation(body) => write!(f, "Video generation violation: {}", body),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
#[cfg(test)]
|
||||
pub (crate) mod test_utils;
|
||||
|
||||
// User API
|
||||
pub mod creds;
|
||||
pub mod error;
|
||||
pub mod requests;
|
||||
pub mod utils;
|
||||
@@ -0,0 +1,386 @@
|
||||
use crate::creds::seedance2pro_session::Seedance2ProSession;
|
||||
use crate::error::seedance2pro_client_error::Seedance2ProClientError;
|
||||
use crate::error::seedance2pro_error::Seedance2ProError;
|
||||
use crate::error::seedance2pro_generic_api_error::Seedance2ProGenericApiError;
|
||||
use crate::error::seedance2pro_specific_api_error::Seedance2ProSpecificApiError;
|
||||
use crate::requests::generate_video::request_types::*;
|
||||
use crate::utils::common_headers::FIREFOX_USER_AGENT;
|
||||
use log::info;
|
||||
use wreq::Client;
|
||||
use wreq_util::Emulation;
|
||||
|
||||
const RUN_TASK_URL: &str = "https://seedance2-pro.com/api/trpc/workflow.runTask?batch=1";
|
||||
|
||||
// --- Request args ---
|
||||
|
||||
pub struct GenerateVideoArgs<'a> {
|
||||
pub session: &'a Seedance2ProSession,
|
||||
|
||||
pub prompt: String,
|
||||
|
||||
pub resolution: Resolution,
|
||||
|
||||
/// Duration in seconds (4–15).
|
||||
pub duration_seconds: u8,
|
||||
|
||||
pub batch_count: BatchCount,
|
||||
|
||||
/// Optional start frame image URL (keyframe mode).
|
||||
pub start_frame_url: Option<String>,
|
||||
|
||||
/// Optional end frame image URL (keyframe mode).
|
||||
pub end_frame_url: Option<String>,
|
||||
|
||||
/// Optional reference image URLs (reference mode). When present, takes priority over start/end frames.
|
||||
pub reference_image_urls: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
impl GenerateVideoArgs<'_> {
|
||||
/// Estimates the credit cost for this generation request.
|
||||
///
|
||||
/// Pricing rules:
|
||||
/// - 40 credits per second of video (4s = 160, 15s = 600)
|
||||
/// - Resolution has no effect on cost
|
||||
/// - Input mode (text, keyframe, reference) has no effect on cost
|
||||
/// - Batch 1 = 1×, Batch 2 = 2×, Batch 4 = 3× (not 4×)
|
||||
pub fn estimate_credits(&self) -> u32 {
|
||||
let per_video = u32::from(self.duration_seconds) * 40;
|
||||
let batch_multiplier = match self.batch_count {
|
||||
BatchCount::One => 1,
|
||||
BatchCount::Two => 2,
|
||||
BatchCount::Four => 3, // NB: 3x, not 4x.
|
||||
};
|
||||
per_video * batch_multiplier
|
||||
}
|
||||
|
||||
pub fn estimate_cost_in_usd_cents(&self) -> u64 {
|
||||
// 25000 Credits costs $99.99 as of 2026-02-20
|
||||
// 250 for $1.
|
||||
let credits = self.estimate_credits() as f64;
|
||||
let cost = credits / 250.0 * 100.0;
|
||||
cost.round() as u64
|
||||
}
|
||||
}
|
||||
|
||||
// --- Public enums ---
|
||||
|
||||
/// Video resolution / aspect ratio.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum Resolution {
|
||||
/// 16:9 landscape (1280x720)
|
||||
Landscape16x9,
|
||||
/// 9:16 portrait (720x1280)
|
||||
Portrait9x16,
|
||||
/// 1:1 square (720x720)
|
||||
Square1x1,
|
||||
/// 4:3 standard (960x720)
|
||||
Standard4x3,
|
||||
/// 3:4 portrait (720x960)
|
||||
Portrait3x4,
|
||||
}
|
||||
|
||||
impl Resolution {
|
||||
fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Landscape16x9 => "1280x720",
|
||||
Self::Portrait9x16 => "720x1280",
|
||||
Self::Square1x1 => "720x720",
|
||||
Self::Standard4x3 => "960x720",
|
||||
Self::Portrait3x4 => "720x960",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Number of videos to generate in a single request.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum BatchCount {
|
||||
One,
|
||||
Two,
|
||||
Four,
|
||||
}
|
||||
|
||||
impl BatchCount {
|
||||
fn as_u8(&self) -> u8 {
|
||||
match self {
|
||||
Self::One => 1,
|
||||
Self::Two => 2,
|
||||
Self::Four => 4,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- Response ---
|
||||
|
||||
pub struct GenerateVideoResponse {
|
||||
pub task_id: String,
|
||||
|
||||
pub order_id: String,
|
||||
|
||||
/// Present when batch_count > 1.
|
||||
pub task_ids: Option<Vec<String>>,
|
||||
|
||||
/// Present when batch_count > 1.
|
||||
pub order_ids: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
// --- Implementation ---
|
||||
|
||||
pub async fn generate_video(args: GenerateVideoArgs<'_>) -> Result<GenerateVideoResponse, Seedance2ProError> {
|
||||
let is_reference_mode = args.reference_image_urls.as_ref().is_some_and(|urls| !urls.is_empty());
|
||||
|
||||
let video_input_mode = if is_reference_mode { "reference" } else { "keyframe" };
|
||||
|
||||
let uploaded_urls: Option<Vec<String>> = if is_reference_mode {
|
||||
args.reference_image_urls
|
||||
} else {
|
||||
let mut urls = Vec::new();
|
||||
if let Some(url) = args.start_frame_url {
|
||||
urls.push(url);
|
||||
}
|
||||
if let Some(url) = args.end_frame_url {
|
||||
urls.push(url);
|
||||
}
|
||||
if urls.is_empty() { None } else { Some(urls) }
|
||||
};
|
||||
|
||||
let batch_count_value = args.batch_count.as_u8();
|
||||
let batch_count = if batch_count_value > 1 { Some(batch_count_value) } else { None };
|
||||
|
||||
let duration = format!("{}s", args.duration_seconds);
|
||||
|
||||
info!(
|
||||
"Generating video: mode={}, resolution={}, duration={}, batch={}",
|
||||
video_input_mode, args.resolution.as_str(), duration, batch_count_value
|
||||
);
|
||||
|
||||
let request_body = BatchRequest {
|
||||
zero: BatchRequestInner {
|
||||
json: BatchRequestJson {
|
||||
business_type: "wan22-video-generation",
|
||||
api_params: ApiParams {
|
||||
prompt: args.prompt,
|
||||
resolution: args.resolution.as_str().to_string(),
|
||||
mode: "normal",
|
||||
model: "seedance-20",
|
||||
duration,
|
||||
video_input_mode,
|
||||
uploaded_urls,
|
||||
batch_count,
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
let cookie = args.session.cookies.as_str();
|
||||
|
||||
let client = Client::builder()
|
||||
.emulation(Emulation::Firefox143)
|
||||
.build()
|
||||
.map_err(|err| Seedance2ProClientError::WreqClientError(err))?;
|
||||
|
||||
let response = client.post(RUN_TASK_URL)
|
||||
.header("User-Agent", FIREFOX_USER_AGENT)
|
||||
.header("Accept", "*/*")
|
||||
.header("Accept-Language", "en-US,en;q=0.9")
|
||||
.header("Accept-Encoding", "gzip, deflate, br, zstd")
|
||||
.header("Referer", "https://seedance2-pro.com/")
|
||||
.header("Content-Type", "application/json")
|
||||
.header("x-trpc-source", "client")
|
||||
.header("Origin", "https://seedance2-pro.com")
|
||||
.header("Connection", "keep-alive")
|
||||
.header("Cookie", cookie)
|
||||
.header("Sec-Fetch-Dest", "empty")
|
||||
.header("Sec-Fetch-Mode", "cors")
|
||||
.header("Sec-Fetch-Site", "same-origin")
|
||||
.header("Priority", "u=4")
|
||||
.header("TE", "trailers")
|
||||
.json(&request_body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| Seedance2ProGenericApiError::WreqError(err))?;
|
||||
|
||||
let status = response.status();
|
||||
let response_body = response.text()
|
||||
.await
|
||||
.map_err(|err| Seedance2ProGenericApiError::WreqError(err))?;
|
||||
|
||||
info!("Response status: {}, body: {}", status, response_body);
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(Seedance2ProGenericApiError::UncategorizedBadResponseWithStatusAndBody {
|
||||
status_code: status,
|
||||
body: response_body,
|
||||
}.into());
|
||||
}
|
||||
|
||||
let batch_response: Vec<BatchResponseItem> = serde_json::from_str(&response_body)
|
||||
.map_err(|err| Seedance2ProGenericApiError::SerdeResponseParseErrorWithBody(err, response_body.clone()))?;
|
||||
|
||||
let task_data = batch_response
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| Seedance2ProGenericApiError::UncategorizedBadResponse(
|
||||
"Empty batch response array".to_string()
|
||||
))?
|
||||
.result
|
||||
.data
|
||||
.json;
|
||||
|
||||
if task_data.violation_warning {
|
||||
return Err(Seedance2ProSpecificApiError::VideoGenerationViolation(response_body).into());
|
||||
}
|
||||
|
||||
Ok(GenerateVideoResponse {
|
||||
task_id: task_data.task_id,
|
||||
order_id: task_data.order_id,
|
||||
task_ids: task_data.task_ids,
|
||||
order_ids: task_data.order_ids,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::creds::seedance2pro_session::Seedance2ProSession;
|
||||
use crate::test_utils::get_test_cookies::get_test_cookies;
|
||||
use crate::test_utils::setup_test_logging::setup_test_logging;
|
||||
use errors::AnyhowResult;
|
||||
use log::LevelFilter;
|
||||
|
||||
fn dummy_session() -> Seedance2ProSession {
|
||||
Seedance2ProSession::from_cookies_string(String::new())
|
||||
}
|
||||
|
||||
fn args_with(duration_seconds: u8, batch_count: BatchCount) -> GenerateVideoArgs<'static> {
|
||||
// Safety: the dummy session is leaked so the reference is 'static for test purposes.
|
||||
let session = Box::leak(Box::new(dummy_session()));
|
||||
GenerateVideoArgs {
|
||||
session,
|
||||
prompt: String::new(),
|
||||
resolution: Resolution::Square1x1,
|
||||
duration_seconds,
|
||||
batch_count,
|
||||
start_frame_url: None,
|
||||
end_frame_url: None,
|
||||
reference_image_urls: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_credits() {
|
||||
// 40 credits per second, batch 1
|
||||
assert_eq!(args_with(4, BatchCount::One).estimate_credits(), 160);
|
||||
assert_eq!(args_with(5, BatchCount::One).estimate_credits(), 200);
|
||||
assert_eq!(args_with(6, BatchCount::One).estimate_credits(), 240);
|
||||
assert_eq!(args_with(7, BatchCount::One).estimate_credits(), 280);
|
||||
assert_eq!(args_with(15, BatchCount::One).estimate_credits(), 600);
|
||||
|
||||
// Batch 2 = 2×
|
||||
assert_eq!(args_with(4, BatchCount::Two).estimate_credits(), 320);
|
||||
assert_eq!(args_with(5, BatchCount::Two).estimate_credits(), 400);
|
||||
assert_eq!(args_with(15, BatchCount::Two).estimate_credits(), 1200);
|
||||
|
||||
// Batch 4 = 3× (not 4×)
|
||||
assert_eq!(args_with(4, BatchCount::Four).estimate_credits(), 480);
|
||||
assert_eq!(args_with(5, BatchCount::Four).estimate_credits(), 600);
|
||||
assert_eq!(args_with(15, BatchCount::Four).estimate_credits(), 1800);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_cost_usd_cents() {
|
||||
// 40 credits per second, batch 1
|
||||
assert_eq!(args_with(4, BatchCount::One).estimate_cost_in_usd_cents(), 64);
|
||||
assert_eq!(args_with(5, BatchCount::One).estimate_cost_in_usd_cents(), 80);
|
||||
assert_eq!(args_with(6, BatchCount::One).estimate_cost_in_usd_cents(), 96);
|
||||
assert_eq!(args_with(7, BatchCount::One).estimate_cost_in_usd_cents(), 112);
|
||||
assert_eq!(args_with(15, BatchCount::One).estimate_cost_in_usd_cents(), 240);
|
||||
|
||||
// Batch 2 = 2×
|
||||
assert_eq!(args_with(4, BatchCount::Two).estimate_cost_in_usd_cents(), 128);
|
||||
assert_eq!(args_with(5, BatchCount::Two).estimate_cost_in_usd_cents(), 160);
|
||||
assert_eq!(args_with(15, BatchCount::Two).estimate_cost_in_usd_cents(), 480);
|
||||
|
||||
// Batch 4 = 3× (not 4×)
|
||||
assert_eq!(args_with(4, BatchCount::Four).estimate_cost_in_usd_cents(), 192);
|
||||
assert_eq!(args_with(5, BatchCount::Four).estimate_cost_in_usd_cents(), 240);
|
||||
assert_eq!(args_with(15, BatchCount::Four).estimate_cost_in_usd_cents(), 720);
|
||||
}
|
||||
|
||||
fn test_session() -> AnyhowResult<Seedance2ProSession> {
|
||||
let cookies = get_test_cookies()?;
|
||||
Ok(Seedance2ProSession::from_cookies_string(cookies))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // manually test — requires real cookies
|
||||
async fn test_generate_text_to_video() -> AnyhowResult<()> {
|
||||
setup_test_logging(LevelFilter::Trace);
|
||||
let session = test_session()?;
|
||||
let args = GenerateVideoArgs {
|
||||
session: &session,
|
||||
prompt: "A corgi eating a cake in a fancy kitchen.".to_string(),
|
||||
resolution: Resolution::Square1x1,
|
||||
duration_seconds: 5,
|
||||
batch_count: BatchCount::One,
|
||||
start_frame_url: None,
|
||||
end_frame_url: None,
|
||||
reference_image_urls: None,
|
||||
};
|
||||
let result = generate_video(args).await?;
|
||||
println!("Task ID: {}", result.task_id);
|
||||
println!("Order ID: {}", result.order_id);
|
||||
assert!(!result.task_id.is_empty());
|
||||
assert!(!result.order_id.is_empty());
|
||||
assert_eq!(1, 2); // NB: Intentional failure to inspect output.
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // manually test — requires real cookies
|
||||
async fn test_generate_keyframe_video() -> AnyhowResult<()> {
|
||||
setup_test_logging(LevelFilter::Trace);
|
||||
let session = test_session()?;
|
||||
let args = GenerateVideoArgs {
|
||||
session: &session,
|
||||
prompt: "A dog shakes the glasses off its head. The camera pans out as the shiba shakes. The shiba barks.".to_string(),
|
||||
resolution: Resolution::Landscape16x9,
|
||||
duration_seconds: 15,
|
||||
batch_count: BatchCount::One,
|
||||
start_frame_url: Some("https://static.seedance2-pro.com/materials/20260219/1771496300184-fb32e08c.jpg".to_string()),
|
||||
end_frame_url: None,
|
||||
reference_image_urls: None,
|
||||
};
|
||||
let result = generate_video(args).await?;
|
||||
println!("Task ID: {}", result.task_id);
|
||||
println!("Order ID: {}", result.order_id);
|
||||
assert!(!result.task_id.is_empty());
|
||||
assert_eq!(1, 2); // NB: Intentional failure to inspect output.
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // manually test — requires real cookies
|
||||
async fn test_generate_reference_video() -> AnyhowResult<()> {
|
||||
setup_test_logging(LevelFilter::Trace);
|
||||
let session = test_session()?;
|
||||
let args = GenerateVideoArgs {
|
||||
session: &session,
|
||||
prompt: "The dog in @2 is in the office at @1 without the man. The office is dark and moonlight streams in through the windows. Particles of dust gleam in the moon beams. Suddenly, the dog jumps on the desk and barks.".to_string(),
|
||||
resolution: Resolution::Standard4x3,
|
||||
duration_seconds: 10,
|
||||
batch_count: BatchCount::One,
|
||||
start_frame_url: None,
|
||||
end_frame_url: None,
|
||||
reference_image_urls: Some(vec![
|
||||
"https://static.seedance2-pro.com/materials/20260219/1771463564512-b14bfe90.png".to_string(),
|
||||
"https://static.seedance2-pro.com/materials/20260219/1771496300184-fb32e08c.jpg".to_string(),
|
||||
]),
|
||||
};
|
||||
let result = generate_video(args).await?;
|
||||
println!("Task ID: {}", result.task_id);
|
||||
println!("Order ID: {}", result.order_id);
|
||||
assert!(!result.task_id.is_empty());
|
||||
assert_eq!(1, 2); // NB: Intentional failure to inspect output.
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod generate_video;
|
||||
mod request_types;
|
||||
@@ -0,0 +1,60 @@
|
||||
use serde_derive::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub (super) struct BatchRequest {
|
||||
#[serde(rename = "0")]
|
||||
pub zero: BatchRequestInner,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub (super) struct BatchRequestInner {
|
||||
pub json: BatchRequestJson,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub (super) struct BatchRequestJson {
|
||||
#[serde(rename = "businessType")]
|
||||
pub business_type: &'static str,
|
||||
#[serde(rename = "apiParams")]
|
||||
pub api_params: ApiParams,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub (super) struct ApiParams {
|
||||
pub prompt: String,
|
||||
pub resolution: String,
|
||||
pub mode: &'static str,
|
||||
pub model: &'static str,
|
||||
pub duration: String,
|
||||
#[serde(rename = "videoInputMode")]
|
||||
pub video_input_mode: &'static str,
|
||||
#[serde(rename = "uploadedUrls", skip_serializing_if = "Option::is_none")]
|
||||
pub uploaded_urls: Option<Vec<String>>,
|
||||
#[serde(rename = "batchCount", skip_serializing_if = "Option::is_none")]
|
||||
pub batch_count: Option<u8>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub (super) struct BatchResponseItem {
|
||||
pub result: BatchResponseResult,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub (super) struct BatchResponseResult {
|
||||
pub data: BatchResponseData,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub (super) struct BatchResponseData {
|
||||
pub json: TaskResponseJson,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub (super) struct TaskResponseJson {
|
||||
pub task_id: String,
|
||||
pub order_id: String,
|
||||
pub task_ids: Option<Vec<String>>,
|
||||
pub order_ids: Option<Vec<String>>,
|
||||
pub violation_warning: bool,
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
pub mod generate_video;
|
||||
pub mod poll_orders;
|
||||
pub mod prepare_image_upload;
|
||||
pub mod upload_image;
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod poll_orders;
|
||||
mod request_types;
|
||||
@@ -0,0 +1,273 @@
|
||||
use crate::creds::seedance2pro_session::Seedance2ProSession;
|
||||
use crate::error::seedance2pro_client_error::Seedance2ProClientError;
|
||||
use crate::error::seedance2pro_error::Seedance2ProError;
|
||||
use crate::error::seedance2pro_generic_api_error::Seedance2ProGenericApiError;
|
||||
use crate::requests::poll_orders::request_types::*;
|
||||
use crate::utils::common_headers::FIREFOX_USER_AGENT;
|
||||
use log::info;
|
||||
use wreq::Client;
|
||||
use wreq_util::Emulation;
|
||||
|
||||
const GET_ORDERS_BASE_URL: &str = "https://seedance2-pro.com/api/trpc/userOrder.getOrders";
|
||||
|
||||
/// Builds the tRPC `input` JSON for the getOrders endpoint.
|
||||
/// When `cursor` is `Some`, it is included in the JSON payload.
|
||||
fn build_input_json(cursor: Option<u64>) -> String {
|
||||
match cursor {
|
||||
None => r#"{"0":{"json":{"limit":30,"format":null,"direction":"forward"},"meta":{"values":{"format":["undefined"]},"v":1}}}"#.to_string(),
|
||||
Some(c) => format!(
|
||||
r#"{{"0":{{"json":{{"limit":30,"format":null,"cursor":{cursor},"direction":"forward"}},"meta":{{"values":{{"format":["undefined"]}},"v":1}}}}}}"#,
|
||||
cursor = c
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
// --- Args & response ---
|
||||
|
||||
pub struct PollOrdersArgs<'a> {
|
||||
pub session: &'a Seedance2ProSession,
|
||||
|
||||
/// Optional cursor from a previous `PollOrdersResponse::next_cursor`.
|
||||
/// When `None`, the most recent orders are returned.
|
||||
pub cursor: Option<u64>,
|
||||
}
|
||||
|
||||
pub struct PollOrdersResponse {
|
||||
pub orders: Vec<OrderStatus>,
|
||||
|
||||
/// Present when there are more orders to fetch.
|
||||
/// Pass this value as `PollOrdersArgs::cursor` in the next call.
|
||||
pub next_cursor: Option<u64>,
|
||||
}
|
||||
|
||||
// --- Public types ---
|
||||
|
||||
/// The lifecycle status of a video generation task.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum TaskStatus {
|
||||
/// The task is queued and has not started yet.
|
||||
Pending,
|
||||
/// The task is actively being processed.
|
||||
Processing,
|
||||
/// The task finished successfully. `result_url` and `results` will be populated.
|
||||
Completed,
|
||||
/// The task failed. `fail_reason` will contain the reason.
|
||||
Failed,
|
||||
/// An unrecognised status string was returned by the server.
|
||||
Unknown(String),
|
||||
}
|
||||
|
||||
impl TaskStatus {
|
||||
fn from_str(s: &str) -> Self {
|
||||
match s {
|
||||
"PENDING" => Self::Pending,
|
||||
"PROCESSING" => Self::Processing,
|
||||
"COMPLETED" => Self::Completed,
|
||||
"FAILED" => Self::Failed,
|
||||
other => Self::Unknown(other.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_terminal(&self) -> bool {
|
||||
matches!(self, Self::Completed | Self::Failed)
|
||||
}
|
||||
}
|
||||
|
||||
/// A single generated video result attached to an order.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct VideoResult {
|
||||
pub url: String,
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
/// Width / height ratio (e.g. 1.777… for 16:9).
|
||||
pub ratio: f64,
|
||||
}
|
||||
|
||||
/// The status of one order (one video generation task).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct OrderStatus {
|
||||
pub order_id: String,
|
||||
|
||||
pub task_status: TaskStatus,
|
||||
|
||||
/// Top-level result video URL. Populated when `task_status` is `Completed`.
|
||||
pub result_url: Option<String>,
|
||||
|
||||
/// Detailed result entries. Typically one entry per video.
|
||||
pub results: Vec<VideoResult>,
|
||||
|
||||
/// Failure reason. Populated when `task_status` is `Failed`.
|
||||
pub fail_reason: Option<String>,
|
||||
|
||||
/// ISO 8601 creation timestamp (e.g. `"2026-02-19T01:20:50.398Z"`).
|
||||
pub created_at: String,
|
||||
}
|
||||
|
||||
// --- Implementation ---
|
||||
|
||||
pub async fn poll_orders(args: PollOrdersArgs<'_>) -> Result<PollOrdersResponse, Seedance2ProError> {
|
||||
info!("Polling orders (cursor: {:?})...", args.cursor);
|
||||
|
||||
let input_json = build_input_json(args.cursor);
|
||||
|
||||
let client = Client::builder()
|
||||
.emulation(Emulation::Firefox143)
|
||||
.build()
|
||||
.map_err(|err| Seedance2ProClientError::WreqClientError(err))?;
|
||||
|
||||
let cookie = args.session.cookies.as_str();
|
||||
|
||||
let request = client.get(GET_ORDERS_BASE_URL)
|
||||
.query(&[("batch", "1"), ("input", input_json.as_str())])
|
||||
.header("User-Agent", FIREFOX_USER_AGENT)
|
||||
.header("Accept", "*/*")
|
||||
.header("Accept-Language", "en-US,en;q=0.9")
|
||||
.header("Accept-Encoding", "gzip, deflate, br, zstd")
|
||||
.header("Referer", "https://seedance2-pro.com/app/gallery")
|
||||
.header("content-type", "application/json")
|
||||
.header("x-trpc-source", "client")
|
||||
.header("Connection", "keep-alive")
|
||||
.header("Cookie", cookie)
|
||||
.header("Sec-Fetch-Dest", "empty")
|
||||
.header("Sec-Fetch-Mode", "cors")
|
||||
.header("Sec-Fetch-Site", "same-origin")
|
||||
.header("Priority", "u=4")
|
||||
.header("TE", "trailers")
|
||||
.build()
|
||||
.map_err(|err| Seedance2ProClientError::WreqClientError(err))?;
|
||||
|
||||
let response = client.execute(request)
|
||||
.await
|
||||
.map_err(|err| Seedance2ProGenericApiError::WreqError(err))?;
|
||||
|
||||
let status = response.status();
|
||||
let response_body = response.text()
|
||||
.await
|
||||
.map_err(|err| Seedance2ProGenericApiError::WreqError(err))?;
|
||||
|
||||
info!("Poll orders response status: {}", status);
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(Seedance2ProGenericApiError::UncategorizedBadResponseWithStatusAndBody {
|
||||
status_code: status,
|
||||
body: response_body,
|
||||
}.into());
|
||||
}
|
||||
|
||||
let batch_response: Vec<BatchResponseItem> = serde_json::from_str(&response_body)
|
||||
.map_err(|err| Seedance2ProGenericApiError::SerdeResponseParseErrorWithBody(err, response_body.clone()))?;
|
||||
|
||||
let json = batch_response
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| Seedance2ProGenericApiError::UncategorizedBadResponse(
|
||||
"Empty batch response array".to_string()
|
||||
))?
|
||||
.result
|
||||
.data
|
||||
.json;
|
||||
|
||||
let next_cursor = json.next_cursor;
|
||||
|
||||
let orders: Vec<OrderStatus> = json.orders
|
||||
.into_iter()
|
||||
.map(|o| OrderStatus {
|
||||
order_id: o.order_id,
|
||||
task_status: TaskStatus::from_str(&o.task_status),
|
||||
result_url: o.result_url,
|
||||
results: o.results.into_iter().map(|r| VideoResult {
|
||||
url: r.url,
|
||||
width: r.width,
|
||||
height: r.height,
|
||||
ratio: r.ratio,
|
||||
}).collect(),
|
||||
fail_reason: o.fail_reason,
|
||||
created_at: o.created_at,
|
||||
})
|
||||
.collect();
|
||||
|
||||
info!("Polled {} orders, next_cursor: {:?}", orders.len(), next_cursor);
|
||||
|
||||
Ok(PollOrdersResponse { orders, next_cursor })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::creds::seedance2pro_session::Seedance2ProSession;
|
||||
use crate::test_utils::get_test_cookies::get_test_cookies;
|
||||
use crate::test_utils::setup_test_logging::setup_test_logging;
|
||||
use errors::AnyhowResult;
|
||||
use log::LevelFilter;
|
||||
|
||||
fn test_session() -> AnyhowResult<Seedance2ProSession> {
|
||||
let cookies = get_test_cookies()?;
|
||||
Ok(Seedance2ProSession::from_cookies_string(cookies))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // manually test — requires real cookies
|
||||
async fn test_poll_all_orders() -> AnyhowResult<()> {
|
||||
setup_test_logging(LevelFilter::Trace);
|
||||
let session = test_session()?;
|
||||
let result = poll_orders(PollOrdersArgs { session: &session, cursor: None }).await?;
|
||||
println!("Orders returned: {}", result.orders.len());
|
||||
println!("Next cursor: {:?}", result.next_cursor);
|
||||
for order in &result.orders {
|
||||
println!(" {} | {:?} | result_url={:?} | fail={:?}",
|
||||
order.order_id, order.task_status, order.result_url, order.fail_reason);
|
||||
}
|
||||
assert_eq!(1, 2); // NB: Intentional failure to inspect output.
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // manually test — requires real cookies and a known cursor value
|
||||
async fn test_poll_with_cursor() -> AnyhowResult<()> {
|
||||
setup_test_logging(LevelFilter::Trace);
|
||||
let session = test_session()?;
|
||||
// Use the cursor value returned from a prior call (e.g. 394062 from the example responses).
|
||||
let cursor: u64 = 394062;
|
||||
let result = poll_orders(PollOrdersArgs { session: &session, cursor: Some(cursor) }).await?;
|
||||
println!("Orders returned (with cursor {}): {}", cursor, result.orders.len());
|
||||
println!("Next cursor: {:?}", result.next_cursor);
|
||||
for order in &result.orders {
|
||||
println!(" {} | {:?} | result_url={:?} | fail={:?}",
|
||||
order.order_id, order.task_status, order.result_url, order.fail_reason);
|
||||
}
|
||||
assert_eq!(1, 2); // NB: Intentional failure to inspect output.
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // manually test — requires real cookies; exhausts all pages
|
||||
async fn test_poll_all_pages() -> AnyhowResult<()> {
|
||||
setup_test_logging(LevelFilter::Trace);
|
||||
let session = test_session()?;
|
||||
|
||||
let mut cursor: Option<u64> = None;
|
||||
let mut page = 0usize;
|
||||
let mut total_orders = 0usize;
|
||||
|
||||
loop {
|
||||
page += 1;
|
||||
let result = poll_orders(PollOrdersArgs { session: &session, cursor }).await?;
|
||||
let page_count = result.orders.len();
|
||||
total_orders += page_count;
|
||||
|
||||
println!("Page {}: {} orders, next_cursor: {:?}", page, page_count, result.next_cursor);
|
||||
for order in &result.orders {
|
||||
println!(" {} | {:?}", order.order_id, order.task_status);
|
||||
}
|
||||
|
||||
cursor = result.next_cursor;
|
||||
if cursor.is_none() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
println!("Total orders across {} pages: {}", page, total_orders);
|
||||
assert_eq!(1, 2); // NB: Intentional failure to inspect output.
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
use serde_derive::Deserialize;
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub(super) struct BatchResponseItem {
|
||||
pub result: BatchResponseResult,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub(super) struct BatchResponseResult {
|
||||
pub data: BatchResponseData,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub(super) struct BatchResponseData {
|
||||
pub json: OrdersResponseJson,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(super) struct OrdersResponseJson {
|
||||
pub orders: Vec<RawOrder>,
|
||||
pub next_cursor: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(super) struct RawOrder {
|
||||
pub order_id: String,
|
||||
pub result_url: Option<String>,
|
||||
pub task_status: String,
|
||||
pub results: Vec<RawVideoResult>,
|
||||
pub fail_reason: Option<String>,
|
||||
pub created_at: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub(super) struct RawVideoResult {
|
||||
pub url: String,
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
pub ratio: f64,
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod prepare_image_upload;
|
||||
mod request_types;
|
||||
+140
@@ -0,0 +1,140 @@
|
||||
use crate::creds::seedance2pro_session::Seedance2ProSession;
|
||||
use crate::error::seedance2pro_client_error::Seedance2ProClientError;
|
||||
use crate::error::seedance2pro_error::Seedance2ProError;
|
||||
use crate::error::seedance2pro_generic_api_error::Seedance2ProGenericApiError;
|
||||
use crate::requests::prepare_image_upload::request_types::*;
|
||||
use crate::utils::common_headers::FIREFOX_USER_AGENT;
|
||||
use chrono::Utc;
|
||||
use log::info;
|
||||
use rand::Rng;
|
||||
use wreq::Client;
|
||||
use wreq_util::Emulation;
|
||||
|
||||
const SIGNED_UPLOAD_URL: &str = "https://seedance2-pro.com/api/trpc/uploads.signedUploadUrl?batch=1";
|
||||
|
||||
/// Generates a material path based on the current time.
|
||||
/// Format: `materials/YYYYMMDD/<unix_millis>-<random_hex>.png`
|
||||
fn generate_material_path() -> String {
|
||||
let now = Utc::now();
|
||||
let date_part = now.format("%Y%m%d").to_string();
|
||||
let timestamp_millis = now.timestamp_millis();
|
||||
let random_hex: u32 = rand::rng().random();
|
||||
let hex_part = format!("{:08x}", random_hex);
|
||||
format!("materials/{}/{}-{}.png", date_part, timestamp_millis, hex_part)
|
||||
}
|
||||
|
||||
pub struct PrepareImageUploadArgs<'a> {
|
||||
pub session: &'a Seedance2ProSession,
|
||||
}
|
||||
|
||||
pub struct PrepareImageUploadResponse {
|
||||
/// The signed URL to upload the image to (Cloudflare R2 / S3-compatible).
|
||||
pub upload_url: String,
|
||||
|
||||
/// The material path that was generated for this upload.
|
||||
pub material_path: String,
|
||||
}
|
||||
|
||||
pub async fn prepare_image_upload(args: PrepareImageUploadArgs<'_>) -> Result<PrepareImageUploadResponse, Seedance2ProError> {
|
||||
let material_path = generate_material_path();
|
||||
|
||||
info!("Preparing image upload with path: {}", material_path);
|
||||
|
||||
let client = Client::builder()
|
||||
.emulation(Emulation::Firefox143)
|
||||
.build()
|
||||
.map_err(|err| Seedance2ProClientError::WreqClientError(err))?;
|
||||
|
||||
let request_body = BatchRequest {
|
||||
zero: BatchRequestInner {
|
||||
json: BatchRequestJson {
|
||||
path: material_path.clone(),
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
let cookie = args.session.cookies.as_str();
|
||||
|
||||
let response = client.post(SIGNED_UPLOAD_URL)
|
||||
.header("User-Agent", FIREFOX_USER_AGENT)
|
||||
.header("Accept", "*/*")
|
||||
.header("Accept-Language", "en-US,en;q=0.9")
|
||||
.header("Accept-Encoding", "gzip, deflate, br, zstd")
|
||||
.header("Referer", "https://seedance2-pro.com/")
|
||||
.header("Content-Type", "application/json")
|
||||
.header("x-trpc-source", "client")
|
||||
.header("Origin", "https://seedance2-pro.com")
|
||||
.header("Connection", "keep-alive")
|
||||
.header("Cookie", cookie)
|
||||
.header("Sec-Fetch-Dest", "empty")
|
||||
.header("Sec-Fetch-Mode", "cors")
|
||||
.header("Sec-Fetch-Site", "same-origin")
|
||||
.header("Priority", "u=4")
|
||||
.header("TE", "trailers")
|
||||
.json(&request_body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| Seedance2ProGenericApiError::WreqError(err))?;
|
||||
|
||||
let status = response.status();
|
||||
let response_body = response.text()
|
||||
.await
|
||||
.map_err(|err| Seedance2ProGenericApiError::WreqError(err))?;
|
||||
|
||||
info!("Response status: {}, body: {}", status, response_body);
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(Seedance2ProGenericApiError::UncategorizedBadResponseWithStatusAndBody {
|
||||
status_code: status,
|
||||
body: response_body,
|
||||
}.into());
|
||||
}
|
||||
|
||||
let batch_response: Vec<BatchResponseItem> = serde_json::from_str(&response_body)
|
||||
.map_err(|err| Seedance2ProGenericApiError::SerdeResponseParseErrorWithBody(err, response_body.clone()))?;
|
||||
|
||||
let upload_url = batch_response
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| Seedance2ProGenericApiError::UncategorizedBadResponse(
|
||||
"Empty batch response array".to_string()
|
||||
))?
|
||||
.result
|
||||
.data
|
||||
.json;
|
||||
|
||||
Ok(PrepareImageUploadResponse {
|
||||
upload_url,
|
||||
material_path,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::creds::seedance2pro_session::Seedance2ProSession;
|
||||
use crate::test_utils::get_test_cookies::get_test_cookies;
|
||||
use crate::test_utils::setup_test_logging::setup_test_logging;
|
||||
use errors::AnyhowResult;
|
||||
use log::LevelFilter;
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // manually test — requires real cookies
|
||||
async fn test_prepare_image_upload() -> AnyhowResult<()> {
|
||||
setup_test_logging(LevelFilter::Trace);
|
||||
let cookies = get_test_cookies()?;
|
||||
let session = Seedance2ProSession::from_cookies_string(cookies);
|
||||
let args = PrepareImageUploadArgs {
|
||||
session: &session,
|
||||
};
|
||||
let result = prepare_image_upload(args).await?;
|
||||
println!("Upload URL: {}", result.upload_url);
|
||||
println!("Material path: {}", result.material_path);
|
||||
assert!(!result.upload_url.is_empty());
|
||||
assert!(result.upload_url.contains("cloudflarestorage.com"));
|
||||
assert!(result.material_path.starts_with("materials/"));
|
||||
|
||||
assert_eq!(1, 2); // NB: Intentional failure to check the response.
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
use serde_derive::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub(super) struct BatchRequest {
|
||||
#[serde(rename = "0")]
|
||||
pub zero: BatchRequestInner,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub(super) struct BatchRequestInner {
|
||||
pub json: BatchRequestJson,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub(super) struct BatchRequestJson {
|
||||
pub path: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub(super) struct BatchResponseItem {
|
||||
pub result: BatchResponseResult,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub(super) struct BatchResponseResult {
|
||||
pub data: BatchResponseData,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub(super) struct BatchResponseData {
|
||||
pub json: String,
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod upload_image;
|
||||
@@ -0,0 +1,127 @@
|
||||
use crate::error::seedance2pro_client_error::Seedance2ProClientError;
|
||||
use crate::error::seedance2pro_error::Seedance2ProError;
|
||||
use crate::error::seedance2pro_generic_api_error::Seedance2ProGenericApiError;
|
||||
use crate::utils::common_headers::FIREFOX_USER_AGENT;
|
||||
use log::info;
|
||||
use url::Url;
|
||||
use wreq::Client;
|
||||
use wreq_util::Emulation;
|
||||
|
||||
const STATIC_BASE_URL: &str = "https://static.seedance2-pro.com";
|
||||
|
||||
pub struct UploadImageArgs {
|
||||
/// The signed upload URL returned by `prepare_image_upload`.
|
||||
pub upload_url: String,
|
||||
|
||||
/// The raw image bytes to upload.
|
||||
pub image_bytes: Vec<u8>,
|
||||
}
|
||||
|
||||
pub struct UploadImageResponse {
|
||||
/// The public-facing URL for the uploaded image.
|
||||
/// e.g. `https://static.seedance2-pro.com/materials/20260219/1771463564512-b14bfe90.png`
|
||||
pub public_url: String,
|
||||
}
|
||||
|
||||
/// Extracts the path from the R2 upload URL and builds the static public URL.
|
||||
/// e.g. `https://comm.….r2.cloudflarestorage.com/materials/20260219/…?X-Amz-…`
|
||||
/// -> `https://static.seedance2-pro.com/materials/20260219/…`
|
||||
fn build_public_url(upload_url: &str) -> Result<String, Seedance2ProError> {
|
||||
let parsed = Url::parse(upload_url)
|
||||
.map_err(|err| Seedance2ProClientError::UrlParseError(err))?;
|
||||
let path = parsed.path(); // e.g. "/materials/20260219/1771463564512-b14bfe90.png"
|
||||
Ok(format!("{}{}", STATIC_BASE_URL, path))
|
||||
}
|
||||
|
||||
pub async fn upload_image(args: UploadImageArgs) -> Result<UploadImageResponse, Seedance2ProError> {
|
||||
info!("Uploading image to: {}", args.upload_url);
|
||||
|
||||
let client = Client::builder()
|
||||
.emulation(Emulation::Firefox143)
|
||||
.build()
|
||||
.map_err(|err| Seedance2ProClientError::WreqClientError(err))?;
|
||||
|
||||
let response = client.put(&args.upload_url)
|
||||
.header("User-Agent", FIREFOX_USER_AGENT)
|
||||
.header("Accept", "*/*")
|
||||
.header("Accept-Language", "en-US,en;q=0.9")
|
||||
.header("Accept-Encoding", "gzip, deflate, br, zstd")
|
||||
.header("Referer", "https://seedance2-pro.com/")
|
||||
.header("Origin", "https://seedance2-pro.com")
|
||||
.header("Connection", "keep-alive")
|
||||
.header("Sec-Fetch-Dest", "empty")
|
||||
.header("Sec-Fetch-Mode", "cors")
|
||||
.header("Sec-Fetch-Site", "cross-site")
|
||||
.header("Priority", "u=4")
|
||||
.body(args.image_bytes)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| Seedance2ProGenericApiError::WreqError(err))?;
|
||||
|
||||
let status = response.status();
|
||||
|
||||
info!("Upload response status: {}", status);
|
||||
|
||||
if !status.is_success() {
|
||||
let body = response.text()
|
||||
.await
|
||||
.map_err(|err| Seedance2ProGenericApiError::WreqError(err))?;
|
||||
|
||||
return Err(Seedance2ProGenericApiError::UncategorizedBadResponseWithStatusAndBody {
|
||||
status_code: status,
|
||||
body,
|
||||
}.into());
|
||||
}
|
||||
|
||||
let public_url = build_public_url(&args.upload_url)?;
|
||||
|
||||
info!("Public URL: {}", public_url);
|
||||
|
||||
Ok(UploadImageResponse { public_url })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::creds::seedance2pro_session::Seedance2ProSession;
|
||||
use crate::test_utils::get_test_cookies::get_test_cookies;
|
||||
use crate::test_utils::setup_test_logging::setup_test_logging;
|
||||
use crate::requests::prepare_image_upload::prepare_image_upload::{
|
||||
prepare_image_upload, PrepareImageUploadArgs,
|
||||
};
|
||||
use errors::AnyhowResult;
|
||||
use log::LevelFilter;
|
||||
use std::fs;
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // manually test — requires real cookies and a test image
|
||||
async fn test_upload_image() -> AnyhowResult<()> {
|
||||
setup_test_logging(LevelFilter::Trace);
|
||||
|
||||
// Step 1: Get a signed upload URL
|
||||
let cookies = get_test_cookies()?;
|
||||
let session = Seedance2ProSession::from_cookies_string(cookies);
|
||||
let prepare_args = PrepareImageUploadArgs {
|
||||
session: &session,
|
||||
};
|
||||
let prepare_result = prepare_image_upload(prepare_args).await?;
|
||||
println!("Upload URL: {}", prepare_result.upload_url);
|
||||
|
||||
// Step 2: Read a test image
|
||||
let image_bytes = fs::read("/Users/bt/dev/storyteller/artcraft/test_data/image/juno.jpg")?;
|
||||
println!("Image size: {} bytes", image_bytes.len());
|
||||
|
||||
// Step 3: Upload
|
||||
let upload_args = UploadImageArgs {
|
||||
upload_url: prepare_result.upload_url,
|
||||
image_bytes,
|
||||
};
|
||||
let result = upload_image(upload_args).await?;
|
||||
println!("Public URL: {}", result.public_url);
|
||||
|
||||
assert!(result.public_url.starts_with("https://static.seedance2-pro.com/materials/"));
|
||||
assert_eq!(1, 2); // NB: Intentional failure to check the response.
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
use errors::AnyhowResult;
|
||||
use std::fs::read_to_string;
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn get_test_cookies() -> AnyhowResult<String> {
|
||||
let cookies = read_to_string("/Users/bt/Artcraft/credentials/seedance2pro_cookies.txt")?;
|
||||
let cookies = cookies.trim().to_string();
|
||||
Ok(cookies)
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
#[cfg(test)]
|
||||
pub (crate) mod get_test_cookies;
|
||||
|
||||
#[cfg(test)]
|
||||
pub (crate) mod setup_test_logging;
|
||||
@@ -0,0 +1,52 @@
|
||||
use chrono::Local;
|
||||
use env_logger::Builder;
|
||||
use log::{debug, error, info, trace, warn, LevelFilter};
|
||||
use std::io::Write;
|
||||
|
||||
const RUST_LOG: &str = "RUST_LOG";
|
||||
|
||||
pub fn setup_test_logging(level: LevelFilter) {
|
||||
let level_str = get_level_str(level);
|
||||
|
||||
println!("Previous log level: {:?}", env::var(RUST_LOG));
|
||||
|
||||
let is_unsafe = env::set_var(RUST_LOG, level_str);
|
||||
|
||||
if is_unsafe.is_none() {
|
||||
println!("Unsafely setting Rust Log Level...");
|
||||
unsafe { std::env::set_var(RUST_LOG, level_str) }
|
||||
}
|
||||
|
||||
println!("Log level: {:?}", env::var(RUST_LOG));
|
||||
|
||||
Builder::new()
|
||||
.is_test(true)
|
||||
.format(|buf, record| {
|
||||
writeln!(buf,
|
||||
"{} [{}] - {}",
|
||||
Local::now().format("%H:%M:%S%.6f"),
|
||||
record.level(),
|
||||
record.args()
|
||||
)
|
||||
})
|
||||
.filter(None, level)
|
||||
.filter_level(level)
|
||||
.init();
|
||||
|
||||
trace!("Test trace log");
|
||||
debug!("Test debug log");
|
||||
info!("Test info log");
|
||||
warn!("Test warn log");
|
||||
error!("Test error log");
|
||||
}
|
||||
|
||||
fn get_level_str(level: LevelFilter) -> &'static str {
|
||||
match level {
|
||||
LevelFilter::Error => "error",
|
||||
LevelFilter::Warn => "warn",
|
||||
LevelFilter::Info => "info",
|
||||
LevelFilter::Debug => "debug",
|
||||
LevelFilter::Trace => "trace",
|
||||
LevelFilter::Off => "off",
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub const FIREFOX_USER_AGENT: &str =
|
||||
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:147.0) Gecko/20100101 Firefox/147.0";
|
||||
@@ -0,0 +1 @@
|
||||
pub mod common_headers;
|
||||
@@ -12,8 +12,10 @@ normal = ["workspace-build-acceleration"]
|
||||
# Workspace hack (faster compile times)
|
||||
workspace-build-acceleration.workspace = true
|
||||
|
||||
# External
|
||||
# Internal
|
||||
easyenv = { path = "../../lib/easyenv" }
|
||||
|
||||
# External
|
||||
log.workspace = true
|
||||
reqwest = { version = "0.12.4", default-features = false, features = [ "charset", "http2", "system-proxy", "rustls-tls" ] }
|
||||
serde_json.workspace = true
|
||||
|
||||
+62
@@ -0,0 +1,62 @@
|
||||
use crate::core::commands::enqueue::generate_error::GenerateError;
|
||||
use crate::core::commands::enqueue::image_to_video::enqueue_image_to_video_command::{EnqueueImageToVideoRequest, SoraOrientation};
|
||||
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
|
||||
use crate::core::events::generation_events::common::{GenerationModel};
|
||||
use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
|
||||
use artcraft_api_defs::generate::video::multi_function::seedance_2p0_multi_function_video_gen::{Seedance2p0AspectRatio, Seedance2p0MultiFunctionVideoGenRequest};
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::tauri::tasks::task_type::TaskType;
|
||||
use idempotency::uuid::generate_random_uuid;
|
||||
use log::{error, info};
|
||||
use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
|
||||
use storyteller_client::endpoints::generate::video::multi_function::seedance_2p0_multi_function_video_gen::seedance_2p0_multi_function_video_gen;
|
||||
|
||||
pub (super) async fn handle_artcraft_seedance_2p0(
|
||||
request: &EnqueueImageToVideoRequest,
|
||||
app_env_configs: &AppEnvConfigs,
|
||||
creds: &StorytellerCredentialSet,
|
||||
) -> Result<TaskEnqueueSuccess, GenerateError> {
|
||||
|
||||
let uuid_idempotency_token = generate_random_uuid();
|
||||
|
||||
let aspect_ratio = match request.sora_orientation {
|
||||
Some(SoraOrientation::Portrait) => Seedance2p0AspectRatio::Portrait9x16,
|
||||
Some(SoraOrientation::Landscape) => Seedance2p0AspectRatio::Landscape16x9,
|
||||
None => Seedance2p0AspectRatio::Landscape16x9,
|
||||
};
|
||||
|
||||
let request = Seedance2p0MultiFunctionVideoGenRequest {
|
||||
uuid_idempotency_token,
|
||||
prompt: request.prompt.clone(),
|
||||
start_frame_media_token: request.image_media_token.clone(),
|
||||
end_frame_media_token: request.end_frame_image_media_token.clone(),
|
||||
reference_image_media_tokens: None,
|
||||
aspect_ratio: Some(aspect_ratio),
|
||||
duration_seconds: None, // TODO: Parameterize
|
||||
batch_count: None,
|
||||
};
|
||||
|
||||
let result = seedance_2p0_multi_function_video_gen(
|
||||
&app_env_configs.storyteller_host,
|
||||
Some(&creds),
|
||||
request,
|
||||
).await;
|
||||
|
||||
let job_token = match result {
|
||||
Ok(enqueued) => {
|
||||
info!("Successfully enqueued.");
|
||||
enqueued.inference_job_token
|
||||
}
|
||||
Err(err) => {
|
||||
error!("Failed to enqueue: {:?}", err);
|
||||
return Err(GenerateError::from(err));
|
||||
}
|
||||
};
|
||||
|
||||
Ok(TaskEnqueueSuccess {
|
||||
task_type: TaskType::VideoGeneration,
|
||||
model: Some(GenerationModel::Seedance2p0),
|
||||
provider: GenerationProvider::Artcraft,
|
||||
provider_job_id: Some(job_token.to_string()),
|
||||
})
|
||||
}
|
||||
+2
@@ -6,6 +6,7 @@ use crate::core::commands::enqueue::image_to_video::artcraft::handle_artcraft_kl
|
||||
use crate::core::commands::enqueue::image_to_video::artcraft::handle_artcraft_kling_2p5_turbo_pro::handle_artcraft_kling_2p5_turbo_pro;
|
||||
use crate::core::commands::enqueue::image_to_video::artcraft::handle_artcraft_kling_2p6_pro::handle_artcraft_kling_2p6_pro;
|
||||
use crate::core::commands::enqueue::image_to_video::artcraft::handle_artcraft_seedance_1_lite::handle_artcraft_seedance_1_lite;
|
||||
use crate::core::commands::enqueue::image_to_video::artcraft::handle_artcraft_seedance_2p0::handle_artcraft_seedance_2p0;
|
||||
use crate::core::commands::enqueue::image_to_video::artcraft::handle_artcraft_sora2::handle_artcraft_sora2;
|
||||
use crate::core::commands::enqueue::image_to_video::artcraft::handle_artcraft_sora2_pro::handle_artcraft_sora2_pro;
|
||||
use crate::core::commands::enqueue::image_to_video::artcraft::handle_artcraft_veo2::handle_artcraft_veo2;
|
||||
@@ -39,6 +40,7 @@ pub async fn handle_video_artcraft(
|
||||
Some(VideoModel::Kling2p5TurboPro) => handle_artcraft_kling_2p5_turbo_pro(request, app_env_configs, &creds).await,
|
||||
Some(VideoModel::Kling2p6Pro) => handle_artcraft_kling_2p6_pro(request, app_env_configs, &creds).await,
|
||||
Some(VideoModel::Seedance10Lite) => handle_artcraft_seedance_1_lite(request, app_env_configs, &creds).await,
|
||||
Some(VideoModel::Seedance2p0) => handle_artcraft_seedance_2p0(request, app_env_configs, &creds).await,
|
||||
Some(VideoModel::Sora2) => handle_artcraft_sora2(request, app_env_configs, &creds).await,
|
||||
Some(VideoModel::Sora2Pro) => handle_artcraft_sora2_pro(request, app_env_configs, &creds).await,
|
||||
Some(VideoModel::Veo2) => handle_artcraft_veo2(request, app_env_configs, &creds).await,
|
||||
|
||||
@@ -5,6 +5,7 @@ mod handle_artcraft_kling_2p1_pro;
|
||||
mod handle_artcraft_kling_2p5_turbo_pro;
|
||||
mod handle_artcraft_kling_2p6_pro;
|
||||
mod handle_artcraft_seedance_1_lite;
|
||||
mod handle_artcraft_seedance_2p0;
|
||||
mod handle_artcraft_sora2;
|
||||
mod handle_artcraft_sora2_pro;
|
||||
mod handle_artcraft_veo2;
|
||||
|
||||
+3
@@ -56,6 +56,9 @@ pub enum VideoModel {
|
||||
#[serde(rename = "seedance_1.0_lite")]
|
||||
Seedance10Lite,
|
||||
|
||||
#[serde(rename = "seedance_2p0")]
|
||||
Seedance2p0,
|
||||
|
||||
#[serde(rename = "sora_2")]
|
||||
Sora2,
|
||||
|
||||
|
||||
@@ -82,6 +82,7 @@ impl TaskEnqueueSuccess{
|
||||
Some(GenerationModel::Kling2p5TurboPro) => Some(TaskModelType::Kling2p5TurboPro),
|
||||
Some(GenerationModel::Kling2p6Pro) => Some(TaskModelType::Kling2p6Pro),
|
||||
Some(GenerationModel::Seedance10Lite) => Some(TaskModelType::Seedance10Lite),
|
||||
Some(GenerationModel::Seedance2p0) => Some(TaskModelType::Seedance2p0),
|
||||
Some(GenerationModel::Sora2) => Some(TaskModelType::Sora2),
|
||||
Some(GenerationModel::Sora2Pro) => Some(TaskModelType::Sora2Pro),
|
||||
Some(GenerationModel::Veo2) => Some(TaskModelType::Veo2),
|
||||
|
||||
@@ -86,6 +86,8 @@ pub enum GenerationModel {
|
||||
Sora2Pro,
|
||||
#[serde(rename = "seedance_1.0_lite")]
|
||||
Seedance10Lite,
|
||||
#[serde(rename = "seedance_2p0")]
|
||||
Seedance2p0,
|
||||
#[serde(rename = "veo_2")]
|
||||
Veo2,
|
||||
#[serde(rename = "veo_3")]
|
||||
|
||||
@@ -16,4 +16,5 @@ pub fn add_artcraft_3d(cors: Cors, _is_production: bool) -> Cors {
|
||||
.allowed_origin("http://localhost:5175") // If already started
|
||||
.allowed_origin("https://macaroni-1.tailce84f.ts.net") // TODO
|
||||
.allowed_origin("https://halide.tailce84f.ts.net") // TODO
|
||||
.allowed_origin("https://brandons-macbook-pro.taild62114.ts.net") // TODO
|
||||
}
|
||||
|
||||
@@ -22,5 +22,8 @@ workspace-build-acceleration.workspace = true
|
||||
once_cell.workspace = true
|
||||
wreq.workspace = true
|
||||
|
||||
[target.'cfg(target_os = "linux")'.dependencies]
|
||||
wreq = { workspace = true, features = ["prefix-symbols"] }
|
||||
|
||||
[dev-dependencies]
|
||||
# (None)
|
||||
|
||||
-1
@@ -86,7 +86,6 @@ pub async fn insert_generic_inference_job_for_fal_queue_with_apriori_job_token<'
|
||||
|
||||
let maybe_external_third_party = InferenceJobExternalThirdParty::Fal;
|
||||
|
||||
let maybe_external_third_party_id = "";
|
||||
const STATUS : JobStatusPlus = JobStatusPlus::Pending;
|
||||
|
||||
let query = sqlx::query!(
|
||||
|
||||
+56
@@ -0,0 +1,56 @@
|
||||
use anyhow::anyhow;
|
||||
use sqlx::MySqlPool;
|
||||
|
||||
use enums::by_table::generic_inference_jobs::frontend_failure_category::FrontendFailureCategory;
|
||||
use enums::common::job_status_plus::JobStatusPlus;
|
||||
use errors::AnyhowResult;
|
||||
use tokens::tokens::generic_inference_jobs::InferenceJobToken;
|
||||
|
||||
pub struct MarkJobFailedByTokenArgs<'a> {
|
||||
pub pool: &'a MySqlPool,
|
||||
pub job_token: &'a InferenceJobToken,
|
||||
pub maybe_public_failure_reason: Option<&'a str>,
|
||||
pub internal_debugging_failure_reason: &'a str,
|
||||
pub maybe_frontend_failure_category: Option<FrontendFailureCategory>,
|
||||
}
|
||||
|
||||
/// Permanently mark an inference job as failed, looked up by its token.
|
||||
/// Unlike `mark_generic_inference_job_failure`, this does not allow retries —
|
||||
/// the job will always land in `complete_failure`.
|
||||
pub async fn mark_job_failed_by_token(args: MarkJobFailedByTokenArgs<'_>) -> AnyhowResult<()> {
|
||||
let maybe_public_failure_reason = args.maybe_public_failure_reason.map(|reason| {
|
||||
let mut reason = reason.trim().to_string();
|
||||
reason.truncate(512);
|
||||
reason
|
||||
});
|
||||
|
||||
let mut internal_debugging_failure_reason = args.internal_debugging_failure_reason.trim().to_string();
|
||||
internal_debugging_failure_reason.truncate(512);
|
||||
|
||||
const FAILURE_STATUS: &str = JobStatusPlus::CompleteFailure.to_str();
|
||||
|
||||
let query_result = sqlx::query!(
|
||||
r#"
|
||||
UPDATE generic_inference_jobs
|
||||
SET
|
||||
status = ?,
|
||||
failure_reason = ?,
|
||||
internal_debugging_failure_reason = ?,
|
||||
frontend_failure_category = ?,
|
||||
retry_at = NULL
|
||||
WHERE token = ?
|
||||
"#,
|
||||
FAILURE_STATUS,
|
||||
maybe_public_failure_reason.as_deref(),
|
||||
&internal_debugging_failure_reason,
|
||||
args.maybe_frontend_failure_category,
|
||||
args.job_token.as_str()
|
||||
)
|
||||
.execute(args.pool)
|
||||
.await;
|
||||
|
||||
match query_result {
|
||||
Err(err) => Err(anyhow!("error with query: {:?}", err)),
|
||||
Ok(_r) => Ok(()),
|
||||
}
|
||||
}
|
||||
@@ -3,6 +3,7 @@ pub mod count_untried_jobs_of_type;
|
||||
pub mod list_available_generic_inference_jobs;
|
||||
pub mod mark_generic_inference_job_completely_failed;
|
||||
pub mod mark_generic_inference_job_failure;
|
||||
pub mod mark_job_failed_by_token;
|
||||
pub mod mark_generic_inference_job_pending_and_grab_lock;
|
||||
pub mod mark_generic_inference_job_successfully_done;
|
||||
pub mod reap_stale_fakeyou_jobs;
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod fal;
|
||||
pub mod job;
|
||||
pub mod seedance2pro;
|
||||
pub mod web;
|
||||
|
||||
+137
@@ -0,0 +1,137 @@
|
||||
use anyhow::anyhow;
|
||||
use log::info;
|
||||
use sqlx::{Executor, MySql};
|
||||
use std::marker::PhantomData;
|
||||
|
||||
use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
|
||||
use enums::by_table::generic_inference_jobs::inference_job_external_third_party::InferenceJobExternalThirdParty;
|
||||
use enums::by_table::generic_inference_jobs::inference_job_product_category::InferenceJobProductCategory;
|
||||
use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType;
|
||||
use enums::common::job_status_plus::JobStatusPlus;
|
||||
use enums::common::visibility::Visibility;
|
||||
use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken;
|
||||
use tokens::tokens::generic_inference_jobs::InferenceJobToken;
|
||||
use tokens::tokens::prompts::PromptToken;
|
||||
use tokens::tokens::users::UserToken;
|
||||
|
||||
use crate::errors::database_query_error::DatabaseQueryError;
|
||||
use crate::payloads::generic_inference_args::generic_inference_args::GenericInferenceArgs;
|
||||
|
||||
pub struct InsertGenericInferenceForSeedance2ProWithAprioriJobTokenArgs<'e, 'c, E>
|
||||
where E: 'e + Executor<'c, Database = MySql>
|
||||
{
|
||||
pub uuid_idempotency_token: &'e str,
|
||||
|
||||
// NOTE: We'll generate this ahead of time so we can save it with billing info!
|
||||
pub apriori_job_token: &'e InferenceJobToken,
|
||||
|
||||
/// The external order_id returned by the seedance2pro generate_video call.
|
||||
pub maybe_external_third_party_id: &'e str,
|
||||
|
||||
pub maybe_inference_args: Option<GenericInferenceArgs>,
|
||||
|
||||
pub maybe_prompt_token: Option<&'e PromptToken>,
|
||||
|
||||
pub maybe_creator_user_token: Option<&'e UserToken>,
|
||||
pub maybe_avt_token: Option<&'e AnonymousVisitorTrackingToken>,
|
||||
pub creator_ip_address: &'e str,
|
||||
pub creator_set_visibility: Visibility,
|
||||
|
||||
pub mysql_executor: E,
|
||||
|
||||
pub phantom: PhantomData<&'c E>,
|
||||
}
|
||||
|
||||
pub async fn insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token<'e, 'c : 'e, E>(
|
||||
args: InsertGenericInferenceForSeedance2ProWithAprioriJobTokenArgs<'e, 'c, E>
|
||||
) -> Result<InferenceJobToken, DatabaseQueryError>
|
||||
where E: 'e + Executor<'c, Database = MySql>
|
||||
{
|
||||
let serialized_args_payload = serde_json::ser::to_string(&args.maybe_inference_args)
|
||||
.map_err(|_e| anyhow!("could not encode inference args"))?;
|
||||
|
||||
const JOB_TYPE: InferenceJobType = InferenceJobType::Seedance2ProQueue;
|
||||
const INFERENCE_CATEGORY: InferenceCategory = InferenceCategory::VideoGeneration;
|
||||
const PRODUCT_CATEGORY: InferenceJobProductCategory = InferenceJobProductCategory::Seedance2ProVideo;
|
||||
const EXTERNAL_THIRD_PARTY: InferenceJobExternalThirdParty = InferenceJobExternalThirdParty::Seedance2Pro;
|
||||
const STATUS: JobStatusPlus = JobStatusPlus::Pending;
|
||||
|
||||
let query = sqlx::query!(
|
||||
r#"
|
||||
INSERT INTO generic_inference_jobs
|
||||
SET
|
||||
token = ?,
|
||||
uuid_idempotency_token = ?,
|
||||
|
||||
job_type = ?,
|
||||
|
||||
maybe_external_third_party = ?,
|
||||
maybe_external_third_party_id = ?,
|
||||
|
||||
product_category = ?,
|
||||
inference_category = ?,
|
||||
|
||||
maybe_model_type = NULL,
|
||||
maybe_model_token = NULL,
|
||||
|
||||
maybe_input_source_token = NULL,
|
||||
maybe_input_source_token_type = NULL,
|
||||
|
||||
maybe_download_url = NULL,
|
||||
maybe_cover_image_media_file_token = NULL,
|
||||
|
||||
maybe_prompt_token = ?,
|
||||
|
||||
maybe_raw_inference_text = NULL,
|
||||
|
||||
maybe_inference_args = ?,
|
||||
|
||||
maybe_creator_user_token = ?,
|
||||
maybe_creator_anonymous_visitor_token = ?,
|
||||
creator_ip_address = ?,
|
||||
creator_set_visibility = ?,
|
||||
|
||||
priority_level = 0,
|
||||
is_keepalive_required = FALSE,
|
||||
max_duration_seconds = 0,
|
||||
|
||||
is_debug_request = FALSE,
|
||||
maybe_routing_tag = NULL,
|
||||
|
||||
status = ?
|
||||
"#,
|
||||
args.apriori_job_token.as_str(),
|
||||
args.uuid_idempotency_token,
|
||||
|
||||
JOB_TYPE.to_str(),
|
||||
|
||||
EXTERNAL_THIRD_PARTY.to_str(),
|
||||
args.maybe_external_third_party_id,
|
||||
|
||||
PRODUCT_CATEGORY.to_str(),
|
||||
INFERENCE_CATEGORY.to_str(),
|
||||
|
||||
args.maybe_prompt_token.map(|t| t.to_string()),
|
||||
|
||||
serialized_args_payload,
|
||||
|
||||
args.maybe_creator_user_token.map(|t| t.to_string()),
|
||||
args.maybe_avt_token.map(|t| t.to_string()),
|
||||
args.creator_ip_address,
|
||||
args.creator_set_visibility.to_str(),
|
||||
|
||||
STATUS.to_str(),
|
||||
);
|
||||
|
||||
let query_result = query.execute(args.mysql_executor)
|
||||
.await;
|
||||
|
||||
let record_id = match query_result {
|
||||
Err(err) => return Err(DatabaseQueryError::from(err)),
|
||||
Ok(res) => res.last_insert_id(),
|
||||
};
|
||||
|
||||
info!("Insert generic inference job for Seedance 2 Pro queue: {} with record ID {}", args.apriori_job_token, record_id);
|
||||
|
||||
Ok(args.apriori_job_token.clone())
|
||||
}
|
||||
+92
@@ -0,0 +1,92 @@
|
||||
use anyhow::anyhow;
|
||||
use log::warn;
|
||||
use sqlx::MySqlPool;
|
||||
|
||||
use enums::by_table::generic_inference_jobs::inference_job_external_third_party::InferenceJobExternalThirdParty;
|
||||
use enums::common::visibility::Visibility;
|
||||
use errors::AnyhowResult;
|
||||
use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken;
|
||||
use tokens::tokens::generic_inference_jobs::InferenceJobToken;
|
||||
use tokens::tokens::prompts::PromptToken;
|
||||
use tokens::tokens::users::UserToken;
|
||||
|
||||
/// A Seedance2Pro job that is waiting for a result from the external API.
|
||||
#[derive(Debug)]
|
||||
pub struct PendingSeedance2ProJob {
|
||||
pub job_token: InferenceJobToken,
|
||||
|
||||
/// The external order_id from seedance2-pro.com.
|
||||
pub order_id: String,
|
||||
|
||||
pub maybe_creator_user_token: Option<UserToken>,
|
||||
pub maybe_creator_anonymous_visitor_token: Option<AnonymousVisitorTrackingToken>,
|
||||
pub creator_ip_address: String,
|
||||
pub creator_set_visibility: Visibility,
|
||||
|
||||
pub maybe_prompt_token: Option<PromptToken>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct RawRecord {
|
||||
job_token: InferenceJobToken,
|
||||
order_id: Option<String>,
|
||||
maybe_creator_user_token: Option<UserToken>,
|
||||
maybe_creator_anonymous_visitor_token: Option<AnonymousVisitorTrackingToken>,
|
||||
creator_ip_address: String,
|
||||
creator_set_visibility: Visibility,
|
||||
maybe_prompt_token: Option<PromptToken>,
|
||||
}
|
||||
|
||||
/// Returns all non-terminal Seedance2Pro jobs that have an associated order_id.
|
||||
pub async fn list_pending_seedance2pro_jobs(pool: &MySqlPool) -> AnyhowResult<Vec<PendingSeedance2ProJob>> {
|
||||
let records = sqlx::query_as!(
|
||||
RawRecord,
|
||||
r#"
|
||||
SELECT
|
||||
jobs.token as `job_token: tokens::tokens::generic_inference_jobs::InferenceJobToken`,
|
||||
jobs.maybe_external_third_party_id as `order_id`,
|
||||
jobs.maybe_creator_user_token as `maybe_creator_user_token: tokens::tokens::users::UserToken`,
|
||||
jobs.maybe_creator_anonymous_visitor_token as `maybe_creator_anonymous_visitor_token: tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken`,
|
||||
jobs.creator_ip_address,
|
||||
jobs.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,
|
||||
jobs.maybe_prompt_token as `maybe_prompt_token: tokens::tokens::prompts::PromptToken`
|
||||
|
||||
FROM generic_inference_jobs as jobs
|
||||
|
||||
WHERE jobs.maybe_external_third_party = ?
|
||||
AND jobs.status NOT IN ('complete_success', 'complete_failure')
|
||||
AND jobs.maybe_external_third_party_id IS NOT NULL
|
||||
|
||||
LIMIT 25000
|
||||
"#,
|
||||
InferenceJobExternalThirdParty::Seedance2Pro.to_str(),
|
||||
)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|err| anyhow!("error querying pending seedance2pro jobs: {:?}", err))?;
|
||||
|
||||
let jobs = records
|
||||
.into_iter()
|
||||
.filter_map(|record| {
|
||||
let order_id = match record.order_id {
|
||||
Some(id) => id,
|
||||
None => {
|
||||
warn!("PendingSeedance2ProJob has no order_id, skipping");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(PendingSeedance2ProJob {
|
||||
job_token: record.job_token,
|
||||
order_id,
|
||||
maybe_creator_user_token: record.maybe_creator_user_token,
|
||||
maybe_creator_anonymous_visitor_token: record.maybe_creator_anonymous_visitor_token,
|
||||
creator_ip_address: record.creator_ip_address,
|
||||
creator_set_visibility: record.creator_set_visibility,
|
||||
maybe_prompt_token: record.maybe_prompt_token,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(jobs)
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token;
|
||||
pub mod list_pending_seedance2pro_jobs;
|
||||
+13
-1
@@ -42,6 +42,11 @@ pub enum FrontendFailureCategory {
|
||||
/// exactly what happened. We'll retry their workload in any case.
|
||||
#[serde(rename = "retryable_worker_error")]
|
||||
RetryableWorkerError,
|
||||
|
||||
/// Model content rules were violated
|
||||
/// Eg. Seedance 2 will report: "your input text violates platform rules. please modify and try again"
|
||||
#[serde(rename = "model_rules_violation")]
|
||||
ModelRulesViolation,
|
||||
}
|
||||
|
||||
// TODO(bt, 2022-12-21): This desperately needs MySQL integration tests!
|
||||
@@ -56,6 +61,7 @@ impl FrontendFailureCategory {
|
||||
Self::KeepAliveElapsed => "keep_alive_elapsed",
|
||||
Self::NotYetImplemented => "not_yet_implemented",
|
||||
Self::RetryableWorkerError => "retryable_worker_error",
|
||||
Self::ModelRulesViolation => "model_rules_violation",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -65,6 +71,7 @@ impl FrontendFailureCategory {
|
||||
"keep_alive_elapsed" => Ok(Self::KeepAliveElapsed),
|
||||
"not_yet_implemented" => Ok(Self::NotYetImplemented),
|
||||
"retryable_worker_error" => Ok(Self::RetryableWorkerError),
|
||||
"model_rules_violation" => Ok(Self::ModelRulesViolation),
|
||||
_ => Err(format!("invalid value: {:?}", value)),
|
||||
}
|
||||
}
|
||||
@@ -77,6 +84,7 @@ impl FrontendFailureCategory {
|
||||
Self::KeepAliveElapsed,
|
||||
Self::NotYetImplemented,
|
||||
Self::RetryableWorkerError,
|
||||
Self::ModelRulesViolation,
|
||||
])
|
||||
}
|
||||
}
|
||||
@@ -95,6 +103,7 @@ mod tests {
|
||||
assert_serialization(FrontendFailureCategory::KeepAliveElapsed, "keep_alive_elapsed");
|
||||
assert_serialization(FrontendFailureCategory::NotYetImplemented, "not_yet_implemented");
|
||||
assert_serialization(FrontendFailureCategory::RetryableWorkerError, "retryable_worker_error");
|
||||
assert_serialization(FrontendFailureCategory::ModelRulesViolation, "model_rules_violation");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -103,6 +112,7 @@ mod tests {
|
||||
assert_eq!(FrontendFailureCategory::KeepAliveElapsed.to_str(), "keep_alive_elapsed");
|
||||
assert_eq!(FrontendFailureCategory::NotYetImplemented.to_str(), "not_yet_implemented");
|
||||
assert_eq!(FrontendFailureCategory::RetryableWorkerError.to_str(), "retryable_worker_error");
|
||||
assert_eq!(FrontendFailureCategory::ModelRulesViolation.to_str(), "model_rules_violation");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -111,16 +121,18 @@ mod tests {
|
||||
assert_eq!(FrontendFailureCategory::from_str("keep_alive_elapsed").unwrap(), FrontendFailureCategory::KeepAliveElapsed);
|
||||
assert_eq!(FrontendFailureCategory::from_str("not_yet_implemented").unwrap(), FrontendFailureCategory::NotYetImplemented);
|
||||
assert_eq!(FrontendFailureCategory::from_str("retryable_worker_error").unwrap(), FrontendFailureCategory::RetryableWorkerError);
|
||||
assert_eq!(FrontendFailureCategory::from_str("model_rules_violation").unwrap(), FrontendFailureCategory::ModelRulesViolation);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_variants() {
|
||||
let mut variants = FrontendFailureCategory::all_variants();
|
||||
assert_eq!(variants.len(), 4);
|
||||
assert_eq!(variants.len(), 5);
|
||||
assert_eq!(variants.pop_first(), Some(FrontendFailureCategory::FaceNotDetected));
|
||||
assert_eq!(variants.pop_first(), Some(FrontendFailureCategory::KeepAliveElapsed));
|
||||
assert_eq!(variants.pop_first(), Some(FrontendFailureCategory::NotYetImplemented));
|
||||
assert_eq!(variants.pop_first(), Some(FrontendFailureCategory::RetryableWorkerError));
|
||||
assert_eq!(variants.pop_first(), Some(FrontendFailureCategory::ModelRulesViolation));
|
||||
assert_eq!(variants.pop_first(), None);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -40,6 +40,7 @@ pub enum InferenceCategory {
|
||||
ImageGeneration,
|
||||
|
||||
/// FAL-powered video generation
|
||||
/// (Also Seedance2-Pro.com)
|
||||
#[serde(rename = "video_generation")]
|
||||
VideoGeneration,
|
||||
|
||||
|
||||
+12
-2
@@ -16,6 +16,10 @@ pub enum InferenceJobExternalThirdParty {
|
||||
#[serde(rename = "fal")]
|
||||
#[default]
|
||||
Fal,
|
||||
|
||||
/// Seedance 2 Pro jobs
|
||||
#[serde(rename = "seedance2pro")]
|
||||
Seedance2Pro,
|
||||
}
|
||||
|
||||
// TODO(bt, 2022-12-21): This desperately needs MySQL integration tests!
|
||||
@@ -27,12 +31,14 @@ impl InferenceJobExternalThirdParty {
|
||||
pub fn to_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Fal => "fal",
|
||||
Self::Seedance2Pro => "seedance2pro",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_str(value: &str) -> Result<Self, String> {
|
||||
match value {
|
||||
"fal" => Ok(Self::Fal),
|
||||
"seedance2pro" => Ok(Self::Seedance2Pro),
|
||||
_ => Err(format!("invalid value: {:?}", value)),
|
||||
}
|
||||
}
|
||||
@@ -42,6 +48,7 @@ impl InferenceJobExternalThirdParty {
|
||||
// NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc.
|
||||
BTreeSet::from([
|
||||
Self::Fal,
|
||||
Self::Seedance2Pro,
|
||||
])
|
||||
}
|
||||
}
|
||||
@@ -57,22 +64,25 @@ mod tests {
|
||||
#[test]
|
||||
fn test_serialization() {
|
||||
assert_serialization(InferenceJobExternalThirdParty::Fal, "fal");
|
||||
assert_serialization(InferenceJobExternalThirdParty::Seedance2Pro, "seedance2pro");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn to_str() {
|
||||
assert_eq!(InferenceJobExternalThirdParty::Fal.to_str(), "fal");
|
||||
assert_eq!(InferenceJobExternalThirdParty::Seedance2Pro.to_str(), "seedance2pro");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_str() {
|
||||
assert_eq!(InferenceJobExternalThirdParty::from_str("fal").unwrap(), InferenceJobExternalThirdParty::Fal);
|
||||
assert_eq!(InferenceJobExternalThirdParty::from_str("seedance2pro").unwrap(), InferenceJobExternalThirdParty::Seedance2Pro);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_variants() {
|
||||
// Static check
|
||||
const EXPECTED_COUNT : usize = 1;
|
||||
const EXPECTED_COUNT : usize = 2;
|
||||
|
||||
assert_eq!(InferenceJobExternalThirdParty::all_variants().len(), EXPECTED_COUNT);
|
||||
assert_eq!(InferenceJobExternalThirdParty::iter().len(), EXPECTED_COUNT);
|
||||
@@ -103,7 +113,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn serialized_length_ok_for_database() {
|
||||
const MAX_LENGTH : usize = 32;
|
||||
const MAX_LENGTH : usize = 16;
|
||||
for variant in InferenceJobExternalThirdParty::all_variants() {
|
||||
let serialized = variant.to_str();
|
||||
assert!(serialized.len() > 0, "variant {:?} is too short", variant);
|
||||
|
||||
+12
-1
@@ -30,6 +30,11 @@ pub enum InferenceJobProductCategory {
|
||||
FalObject,
|
||||
FalBgRemoval,
|
||||
|
||||
// =============== SEEDANCE 2 PRO ===============
|
||||
|
||||
#[serde(rename = "seedance2pro_video")]
|
||||
Seedance2ProVideo,
|
||||
|
||||
// =============== TEXT TO SPEECH ===============
|
||||
|
||||
/// TTS: GptSoVits
|
||||
@@ -125,6 +130,7 @@ impl InferenceJobProductCategory {
|
||||
Self::FalVideo => "fal_video",
|
||||
Self::FalObject => "fal_object",
|
||||
Self::FalBgRemoval => "fal_bg_removal",
|
||||
Self::Seedance2ProVideo => "seedance2pro_video",
|
||||
Self::TtsGptSoVits => "tts_gpt_so_vits",
|
||||
Self::TtsStyleTts2 => "tts_style_tts2",
|
||||
Self::TtsTacotron2 => "tts_tacotron2",
|
||||
@@ -157,6 +163,7 @@ impl InferenceJobProductCategory {
|
||||
"fal_video" => Ok(Self::FalVideo),
|
||||
"fal_object" => Ok(Self::FalObject),
|
||||
"fal_bg_removal" => Ok(Self::FalBgRemoval),
|
||||
"seedance2pro_video" => Ok(Self::Seedance2ProVideo),
|
||||
"tts_gpt_so_vits" => Ok(Self::TtsGptSoVits),
|
||||
"tts_style_tts2" => Ok(Self::TtsStyleTts2),
|
||||
"tts_tacotron2" => Ok(Self::TtsTacotron2),
|
||||
@@ -192,6 +199,7 @@ impl InferenceJobProductCategory {
|
||||
Self::FalVideo,
|
||||
Self::FalObject,
|
||||
Self::FalBgRemoval,
|
||||
Self::Seedance2ProVideo,
|
||||
Self::TtsGptSoVits,
|
||||
Self::TtsStyleTts2,
|
||||
Self::TtsTacotron2,
|
||||
@@ -233,6 +241,7 @@ mod tests {
|
||||
assert_serialization(InferenceJobProductCategory::FalVideo, "fal_video");
|
||||
assert_serialization(InferenceJobProductCategory::FalObject, "fal_object");
|
||||
assert_serialization(InferenceJobProductCategory::FalBgRemoval, "fal_bg_removal");
|
||||
assert_serialization(InferenceJobProductCategory::Seedance2ProVideo, "seedance2pro_video");
|
||||
assert_serialization(InferenceJobProductCategory::TtsGptSoVits, "tts_gpt_so_vits");
|
||||
assert_serialization(InferenceJobProductCategory::TtsStyleTts2, "tts_style_tts2");
|
||||
assert_serialization(InferenceJobProductCategory::TtsTacotron2, "tts_tacotron2");
|
||||
@@ -264,6 +273,7 @@ mod tests {
|
||||
assert_eq!(InferenceJobProductCategory::FalVideo.to_str(), "fal_video");
|
||||
assert_eq!(InferenceJobProductCategory::FalObject.to_str(), "fal_object");
|
||||
assert_eq!(InferenceJobProductCategory::FalBgRemoval.to_str(), "fal_bg_removal");
|
||||
assert_eq!(InferenceJobProductCategory::Seedance2ProVideo.to_str(), "seedance2pro_video");
|
||||
assert_eq!(InferenceJobProductCategory::TtsGptSoVits.to_str(), "tts_gpt_so_vits");
|
||||
assert_eq!(InferenceJobProductCategory::TtsStyleTts2.to_str(), "tts_style_tts2");
|
||||
assert_eq!(InferenceJobProductCategory::TtsTacotron2.to_str(), "tts_tacotron2");
|
||||
@@ -294,6 +304,7 @@ mod tests {
|
||||
assert_eq!(InferenceJobProductCategory::from_str("fal_video").unwrap(), InferenceJobProductCategory::FalVideo);
|
||||
assert_eq!(InferenceJobProductCategory::from_str("fal_object").unwrap(), InferenceJobProductCategory::FalObject);
|
||||
assert_eq!(InferenceJobProductCategory::from_str("fal_bg_removal").unwrap(), InferenceJobProductCategory::FalBgRemoval);
|
||||
assert_eq!(InferenceJobProductCategory::from_str("seedance2pro_video").unwrap(), InferenceJobProductCategory::Seedance2ProVideo);
|
||||
assert_eq!(InferenceJobProductCategory::from_str("tts_gpt_so_vits").unwrap(), InferenceJobProductCategory::TtsGptSoVits);
|
||||
assert_eq!(InferenceJobProductCategory::from_str("tts_style_tts2").unwrap(), InferenceJobProductCategory::TtsStyleTts2);
|
||||
assert_eq!(InferenceJobProductCategory::from_str("tts_tacotron2").unwrap(), InferenceJobProductCategory::TtsTacotron2);
|
||||
@@ -320,7 +331,7 @@ mod tests {
|
||||
#[test]
|
||||
fn all_variants() {
|
||||
// Static check
|
||||
const EXPECTED_COUNT : usize = 27;
|
||||
const EXPECTED_COUNT : usize = 28;
|
||||
|
||||
assert_eq!(InferenceJobProductCategory::all_variants().len(), EXPECTED_COUNT);
|
||||
assert_eq!(InferenceJobProductCategory::iter().len(), EXPECTED_COUNT);
|
||||
|
||||
+12
-2
@@ -18,7 +18,11 @@ pub enum InferenceJobType {
|
||||
/// Fal jobs do not run on our job runners.
|
||||
/// Instead, we receive web hook updates from Fal.
|
||||
FalQueue,
|
||||
|
||||
|
||||
/// Seedance 2 Pro jobs. We poll for results.
|
||||
#[serde(rename = "seedance2pro_queue")]
|
||||
Seedance2ProQueue,
|
||||
|
||||
/// Storyteller Studio and Video Style Transfer Jobs (which we may want to split).
|
||||
/// These run in Comfy.
|
||||
/// TODO(bt,2024-07-15): We may segregate these two job types in the future
|
||||
@@ -110,6 +114,7 @@ impl InferenceJobType {
|
||||
pub fn to_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::FalQueue => "fal_queue",
|
||||
Self::Seedance2ProQueue => "seedance2pro_queue",
|
||||
Self::VideoRender => "video_render",
|
||||
Self::LivePortrait => "live_portrait",
|
||||
Self::FaceFusion => "face_fusion",
|
||||
@@ -136,6 +141,7 @@ impl InferenceJobType {
|
||||
pub fn from_str(value: &str) -> Result<Self, String> {
|
||||
match value {
|
||||
"fal_queue" => Ok(Self::FalQueue),
|
||||
"seedance2pro_queue" => Ok(Self::Seedance2ProQueue),
|
||||
"video_render" => Ok(Self::VideoRender),
|
||||
"live_portrait" => Ok(Self::LivePortrait),
|
||||
"face_fusion" => Ok(Self::FaceFusion),
|
||||
@@ -165,6 +171,7 @@ impl InferenceJobType {
|
||||
// NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc.
|
||||
BTreeSet::from([
|
||||
Self::FalQueue,
|
||||
Self::Seedance2ProQueue,
|
||||
Self::VideoRender,
|
||||
Self::LivePortrait,
|
||||
Self::FaceFusion,
|
||||
@@ -205,6 +212,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_serialization() {
|
||||
assert_serialization(InferenceJobType::FalQueue, "fal_queue");
|
||||
assert_serialization(InferenceJobType::Seedance2ProQueue, "seedance2pro_queue");
|
||||
assert_serialization(InferenceJobType::VideoRender, "video_render");
|
||||
assert_serialization(InferenceJobType::LivePortrait, "live_portrait");
|
||||
assert_serialization(InferenceJobType::FaceFusion, "face_fusion");
|
||||
@@ -230,6 +238,7 @@ mod tests {
|
||||
#[test]
|
||||
fn to_str() {
|
||||
assert_eq!(InferenceJobType::FalQueue.to_str(), "fal_queue");
|
||||
assert_eq!(InferenceJobType::Seedance2ProQueue.to_str(), "seedance2pro_queue");
|
||||
assert_eq!(InferenceJobType::VideoRender.to_str(), "video_render");
|
||||
assert_eq!(InferenceJobType::LivePortrait.to_str(), "live_portrait");
|
||||
assert_eq!(InferenceJobType::FaceFusion.to_str(), "face_fusion");
|
||||
@@ -255,6 +264,7 @@ mod tests {
|
||||
#[test]
|
||||
fn from_str() {
|
||||
assert_eq!(InferenceJobType::from_str("fal_queue").unwrap(), InferenceJobType::FalQueue);
|
||||
assert_eq!(InferenceJobType::from_str("seedance2pro_queue").unwrap(), InferenceJobType::Seedance2ProQueue);
|
||||
assert_eq!(InferenceJobType::from_str("video_render").unwrap(), InferenceJobType::VideoRender);
|
||||
assert_eq!(InferenceJobType::from_str("live_portrait").unwrap(), InferenceJobType::LivePortrait);
|
||||
assert_eq!(InferenceJobType::from_str("face_fusion").unwrap(), InferenceJobType::FaceFusion);
|
||||
@@ -280,7 +290,7 @@ mod tests {
|
||||
#[test]
|
||||
fn all_variants() {
|
||||
// Static check
|
||||
const EXPECTED_COUNT : usize = 21;
|
||||
const EXPECTED_COUNT : usize = 22;
|
||||
|
||||
assert_eq!(InferenceJobType::all_variants().len(), EXPECTED_COUNT);
|
||||
assert_eq!(InferenceJobType::iter().len(), EXPECTED_COUNT);
|
||||
|
||||
@@ -216,7 +216,7 @@ mod tests {
|
||||
#[test]
|
||||
fn all_variants() {
|
||||
let mut variants = MediaFileOriginModelType::all_variants();
|
||||
assert_eq!(variants.len(), 15);
|
||||
assert_eq!(variants.len(), 18);
|
||||
assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::FaceFusion));
|
||||
assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::F5TTS));
|
||||
assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::LivePortrait));
|
||||
|
||||
+11
-1
@@ -66,6 +66,9 @@ pub enum MediaFileOriginProductCategory {
|
||||
#[serde(rename = "image_gen")]
|
||||
ImageGeneration,
|
||||
|
||||
#[serde(rename = "video_gen")]
|
||||
VideoGeneration,
|
||||
|
||||
// Media files for video filters
|
||||
#[deprecated(note = "This isn't relevant product surface area anymore")]
|
||||
#[serde(rename = "video_filter")]
|
||||
@@ -98,6 +101,7 @@ impl MediaFileOriginProductCategory {
|
||||
Self::ZeroShotVoice => "zs_voice",
|
||||
Self::Mocap => "mocap",
|
||||
Self::ImageGeneration => "image_gen",
|
||||
Self::VideoGeneration => "video_gen",
|
||||
Self::VideoFilter => "video_filter",
|
||||
Self::Workflow => "workflow",
|
||||
}
|
||||
@@ -117,6 +121,7 @@ impl MediaFileOriginProductCategory {
|
||||
"zs_voice" => Ok(Self::ZeroShotVoice),
|
||||
"mocap" => Ok(Self::Mocap),
|
||||
"image_gen" => Ok(Self::ImageGeneration),
|
||||
"video_gen" => Ok(Self::VideoGeneration),
|
||||
"video_filter" => Ok(Self::VideoFilter),
|
||||
"workflow" => Ok(Self::Workflow),
|
||||
_ => Err(format!("invalid value: {:?}", value)),
|
||||
@@ -139,6 +144,7 @@ impl MediaFileOriginProductCategory {
|
||||
Self::ZeroShotVoice,
|
||||
Self::Mocap,
|
||||
Self::ImageGeneration,
|
||||
Self::VideoGeneration,
|
||||
Self::VideoFilter,
|
||||
Self::Workflow,
|
||||
])
|
||||
@@ -167,6 +173,7 @@ mod tests {
|
||||
assert_serialization(MediaFileOriginProductCategory::ZeroShotVoice, "zs_voice");
|
||||
assert_serialization(MediaFileOriginProductCategory::Mocap, "mocap");
|
||||
assert_serialization(MediaFileOriginProductCategory::ImageGeneration, "image_gen");
|
||||
assert_serialization(MediaFileOriginProductCategory::VideoGeneration, "video_gen");
|
||||
assert_serialization(MediaFileOriginProductCategory::VideoFilter, "video_filter");
|
||||
assert_serialization(MediaFileOriginProductCategory::Workflow, "workflow");
|
||||
}
|
||||
@@ -185,6 +192,7 @@ mod tests {
|
||||
assert_eq!(MediaFileOriginProductCategory::ZeroShotVoice.to_str(), "zs_voice");
|
||||
assert_eq!(MediaFileOriginProductCategory::Mocap.to_str(), "mocap");
|
||||
assert_eq!(MediaFileOriginProductCategory::ImageGeneration.to_str(), "image_gen");
|
||||
assert_eq!(MediaFileOriginProductCategory::VideoGeneration.to_str(), "video_gen");
|
||||
assert_eq!(MediaFileOriginProductCategory::VideoFilter.to_str(), "video_filter");
|
||||
assert_eq!(MediaFileOriginProductCategory::Workflow.to_str(), "workflow");
|
||||
}
|
||||
@@ -203,6 +211,7 @@ mod tests {
|
||||
assert_eq!(MediaFileOriginProductCategory::from_str("zs_voice").unwrap(), MediaFileOriginProductCategory::ZeroShotVoice);
|
||||
assert_eq!(MediaFileOriginProductCategory::from_str("mocap").unwrap(), MediaFileOriginProductCategory::Mocap);
|
||||
assert_eq!(MediaFileOriginProductCategory::from_str("image_gen").unwrap(),MediaFileOriginProductCategory::ImageGeneration);
|
||||
assert_eq!(MediaFileOriginProductCategory::from_str("video_gen").unwrap(), MediaFileOriginProductCategory::VideoGeneration);
|
||||
assert_eq!(MediaFileOriginProductCategory::from_str("video_filter").unwrap(), MediaFileOriginProductCategory::VideoFilter);
|
||||
assert_eq!(MediaFileOriginProductCategory::from_str("workflow").unwrap(), MediaFileOriginProductCategory::Workflow);
|
||||
}
|
||||
@@ -210,7 +219,7 @@ mod tests {
|
||||
#[test]
|
||||
fn all_variants() {
|
||||
let mut variants = MediaFileOriginProductCategory::all_variants();
|
||||
assert_eq!(variants.len(), 14);
|
||||
assert_eq!(variants.len(), 15);
|
||||
assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Unknown));
|
||||
assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::FaceAnimator));
|
||||
assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::FaceFusion));
|
||||
@@ -223,6 +232,7 @@ mod tests {
|
||||
assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::ZeroShotVoice));
|
||||
assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Mocap));
|
||||
assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::ImageGeneration));
|
||||
assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::VideoGeneration));
|
||||
assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::VideoFilter));
|
||||
assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Workflow));
|
||||
assert_eq!(variants.pop_first(), None);
|
||||
|
||||
@@ -43,7 +43,7 @@ impl Default for JobStatusPlus {
|
||||
}
|
||||
|
||||
impl JobStatusPlus {
|
||||
pub fn to_str(&self) -> &'static str {
|
||||
pub const fn to_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Pending => "pending",
|
||||
Self::Started => "started",
|
||||
|
||||
@@ -93,6 +93,8 @@ pub enum ModelType {
|
||||
Seedance10Lite,
|
||||
#[serde(rename = "seedance_1p0_pro")]
|
||||
Seedance10Pro,
|
||||
#[serde(rename = "seedance_2p0")]
|
||||
Seedance2p0,
|
||||
#[serde(rename = "sora_2")]
|
||||
Sora2,
|
||||
#[serde(rename = "sora_2_pro")]
|
||||
@@ -163,6 +165,7 @@ impl ModelType {
|
||||
Self::Kling2p6Pro => "kling_2p6_pro",
|
||||
Self::Seedance10Lite => "seedance_1p0_lite",
|
||||
Self::Seedance10Pro => "seedance_1p0_pro",
|
||||
Self::Seedance2p0 => "seedance_2p0",
|
||||
Self::Sora2 => "sora_2",
|
||||
Self::Sora2Pro => "sora_2_pro",
|
||||
Self::Veo2 => "veo_2",
|
||||
@@ -217,6 +220,7 @@ impl ModelType {
|
||||
"kling_2p6_pro" => Ok(Self::Kling2p6Pro),
|
||||
"seedance_1p0_lite" => Ok(Self::Seedance10Lite),
|
||||
"seedance_1p0_pro" => Ok(Self::Seedance10Pro),
|
||||
"seedance_2p0" => Ok(Self::Seedance2p0),
|
||||
"sora_2" => Ok(Self::Sora2),
|
||||
"sora_2_pro" => Ok(Self::Sora2Pro),
|
||||
"veo_2" => Ok(Self::Veo2),
|
||||
@@ -275,6 +279,7 @@ impl ModelType {
|
||||
Self::Kling2p6Pro,
|
||||
Self::Seedance10Lite,
|
||||
Self::Seedance10Pro,
|
||||
Self::Seedance2p0,
|
||||
Self::Sora2,
|
||||
Self::Sora2Pro,
|
||||
Self::Veo2,
|
||||
@@ -337,6 +342,7 @@ mod tests {
|
||||
assert_serialization(ModelType::Kling2p6Pro, "kling_2p6_pro");
|
||||
assert_serialization(ModelType::Seedance10Lite, "seedance_1p0_lite");
|
||||
assert_serialization(ModelType::Seedance10Pro, "seedance_1p0_pro");
|
||||
assert_serialization(ModelType::Seedance2p0, "seedance_2p0");
|
||||
assert_serialization(ModelType::Sora2, "sora_2");
|
||||
assert_serialization(ModelType::Sora2Pro, "sora_2_pro");
|
||||
assert_serialization(ModelType::Veo2, "veo_2");
|
||||
@@ -389,6 +395,7 @@ mod tests {
|
||||
assert_eq!(ModelType::Kling2p6Pro.to_str(), "kling_2p6_pro");
|
||||
assert_eq!(ModelType::Seedance10Lite.to_str(), "seedance_1p0_lite");
|
||||
assert_eq!(ModelType::Seedance10Pro.to_str(), "seedance_1p0_pro");
|
||||
assert_eq!(ModelType::Seedance2p0.to_str(), "seedance_2p0");
|
||||
assert_eq!(ModelType::Sora2.to_str(), "sora_2");
|
||||
assert_eq!(ModelType::Sora2Pro.to_str(), "sora_2_pro");
|
||||
assert_eq!(ModelType::Veo2.to_str(), "veo_2");
|
||||
@@ -441,6 +448,7 @@ mod tests {
|
||||
assert_eq!(ModelType::from_str("kling_2p6_pro").unwrap(), ModelType::Kling2p6Pro);
|
||||
assert_eq!(ModelType::from_str("seedance_1p0_lite").unwrap(), ModelType::Seedance10Lite);
|
||||
assert_eq!(ModelType::from_str("seedance_1p0_pro").unwrap(), ModelType::Seedance10Pro);
|
||||
assert_eq!(ModelType::from_str("seedance_2p0").unwrap(), ModelType::Seedance2p0);
|
||||
assert_eq!(ModelType::from_str("sora_2").unwrap(), ModelType::Sora2);
|
||||
assert_eq!(ModelType::from_str("sora_2_pro").unwrap(), ModelType::Sora2Pro);
|
||||
assert_eq!(ModelType::from_str("veo_2").unwrap(), ModelType::Veo2);
|
||||
@@ -458,7 +466,7 @@ mod tests {
|
||||
#[test]
|
||||
fn all_variants() {
|
||||
let mut variants = ModelType::all_variants();
|
||||
assert_eq!(variants.len(), 44);
|
||||
assert_eq!(variants.len(), 45);
|
||||
// Image models
|
||||
assert_eq!(variants.pop_first(), Some(ModelType::Flux1Dev));
|
||||
assert_eq!(variants.pop_first(), Some(ModelType::Flux1Schnell));
|
||||
@@ -495,6 +503,7 @@ mod tests {
|
||||
assert_eq!(variants.pop_first(), Some(ModelType::Kling2p6Pro));
|
||||
assert_eq!(variants.pop_first(), Some(ModelType::Seedance10Lite));
|
||||
assert_eq!(variants.pop_first(), Some(ModelType::Seedance10Pro));
|
||||
assert_eq!(variants.pop_first(), Some(ModelType::Seedance2p0));
|
||||
assert_eq!(variants.pop_first(), Some(ModelType::Sora2));
|
||||
assert_eq!(variants.pop_first(), Some(ModelType::Sora2Pro));
|
||||
assert_eq!(variants.pop_first(), Some(ModelType::Veo2));
|
||||
|
||||
@@ -65,6 +65,8 @@ pub enum TaskModelType {
|
||||
Kling2p6Pro,
|
||||
#[serde(rename = "seedance_1.0_lite")]
|
||||
Seedance10Lite,
|
||||
#[serde(rename = "seedance_2p0")]
|
||||
Seedance2p0,
|
||||
#[serde(rename = "sora_2")]
|
||||
Sora2,
|
||||
#[serde(rename = "sora_2_pro")]
|
||||
@@ -126,6 +128,7 @@ impl TaskModelType {
|
||||
Self::Kling2p5TurboPro => "kling_2p5_turbo_pro",
|
||||
Self::Kling2p6Pro => "kling_2p6_pro",
|
||||
Self::Seedance10Lite => "seedance_1.0_lite",
|
||||
Self::Seedance2p0 => "seedance_2p0",
|
||||
Self::Sora2 => "sora_2",
|
||||
Self::Sora2Pro => "sora_2_pro",
|
||||
Self::Veo2 => "veo_2",
|
||||
@@ -169,6 +172,7 @@ impl TaskModelType {
|
||||
"kling_2p5_turbo_pro" => Ok(Self::Kling2p5TurboPro),
|
||||
"kling_2p6_pro" => Ok(Self::Kling2p6Pro),
|
||||
"seedance_1.0_lite" => Ok(Self::Seedance10Lite),
|
||||
"seedance_2p0" => Ok(Self::Seedance2p0),
|
||||
"sora_2" => Ok(Self::Sora2),
|
||||
"sora_2_pro" => Ok(Self::Sora2Pro),
|
||||
"veo_2" => Ok(Self::Veo2),
|
||||
@@ -215,6 +219,7 @@ impl TaskModelType {
|
||||
Self::Kling2p5TurboPro,
|
||||
Self::Kling2p6Pro,
|
||||
Self::Seedance10Lite,
|
||||
Self::Seedance2p0,
|
||||
Self::Sora2,
|
||||
Self::Sora2Pro,
|
||||
Self::Veo2,
|
||||
@@ -268,6 +273,7 @@ mod tests {
|
||||
assert_serialization(TaskModelType::Kling2p5TurboPro, "kling_2p5_turbo_pro");
|
||||
assert_serialization(TaskModelType::Kling2p6Pro, "kling_2p6_pro");
|
||||
assert_serialization(TaskModelType::Seedance10Lite, "seedance_1.0_lite");
|
||||
assert_serialization(TaskModelType::Seedance2p0, "seedance_2p0");
|
||||
assert_serialization(TaskModelType::Sora2, "sora_2");
|
||||
assert_serialization(TaskModelType::Sora2Pro, "sora_2_pro");
|
||||
assert_serialization(TaskModelType::Veo2, "veo_2");
|
||||
@@ -310,6 +316,7 @@ mod tests {
|
||||
assert_eq!(TaskModelType::Kling2p5TurboPro.to_str(), "kling_2p5_turbo_pro");
|
||||
assert_eq!(TaskModelType::Kling2p6Pro.to_str(), "kling_2p6_pro");
|
||||
assert_eq!(TaskModelType::Seedance10Lite.to_str(), "seedance_1.0_lite");
|
||||
assert_eq!(TaskModelType::Seedance2p0.to_str(), "seedance_2p0");
|
||||
assert_eq!(TaskModelType::Sora2.to_str(), "sora_2");
|
||||
assert_eq!(TaskModelType::Sora2Pro.to_str(), "sora_2_pro");
|
||||
assert_eq!(TaskModelType::Veo2.to_str(), "veo_2");
|
||||
@@ -352,6 +359,7 @@ mod tests {
|
||||
assert_eq!(TaskModelType::from_str("kling_2p5_turbo_pro").unwrap(), TaskModelType::Kling2p5TurboPro);
|
||||
assert_eq!(TaskModelType::from_str("kling_2p6_pro").unwrap(), TaskModelType::Kling2p6Pro);
|
||||
assert_eq!(TaskModelType::from_str("seedance_1.0_lite").unwrap(), TaskModelType::Seedance10Lite);
|
||||
assert_eq!(TaskModelType::from_str("seedance_2p0").unwrap(), TaskModelType::Seedance2p0);
|
||||
assert_eq!(TaskModelType::from_str("sora_2").unwrap(), TaskModelType::Sora2);
|
||||
assert_eq!(TaskModelType::from_str("sora_2_pro").unwrap(), TaskModelType::Sora2Pro);
|
||||
assert_eq!(TaskModelType::from_str("veo_2").unwrap(), TaskModelType::Veo2);
|
||||
@@ -380,7 +388,7 @@ mod tests {
|
||||
#[test]
|
||||
fn all_variants() {
|
||||
let mut variants = TaskModelType::all_variants();
|
||||
assert_eq!(variants.len(), 35);
|
||||
assert_eq!(variants.len(), 36);
|
||||
// Image models
|
||||
assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Dev));
|
||||
assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Schnell));
|
||||
@@ -407,6 +415,7 @@ mod tests {
|
||||
assert_eq!(variants.pop_first(), Some(TaskModelType::Kling2p5TurboPro));
|
||||
assert_eq!(variants.pop_first(), Some(TaskModelType::Kling2p6Pro));
|
||||
assert_eq!(variants.pop_first(), Some(TaskModelType::Seedance10Lite));
|
||||
assert_eq!(variants.pop_first(), Some(TaskModelType::Seedance2p0));
|
||||
assert_eq!(variants.pop_first(), Some(TaskModelType::Sora2));
|
||||
assert_eq!(variants.pop_first(), Some(TaskModelType::Sora2Pro));
|
||||
assert_eq!(variants.pop_first(), Some(TaskModelType::Veo2));
|
||||
|
||||
+6
@@ -100,6 +100,10 @@ pub enum PublicInferenceJobType {
|
||||
/// Sora GPT 4o image gen
|
||||
#[serde(rename = "image_gen_api")]
|
||||
ImageGenApi,
|
||||
|
||||
/// Seedance 2 Pro video generation jobs. We poll for results.
|
||||
#[serde(rename = "seedance2pro_queue")]
|
||||
Seedance2ProQueue,
|
||||
}
|
||||
|
||||
/// NB: Legacy API for older code.
|
||||
@@ -130,6 +134,7 @@ impl PublicInferenceJobType {
|
||||
InferenceJobType::BevyToWorkflow => Self::BevyToWorkflow,
|
||||
InferenceJobType::RerenderAVideo => Self::RerenderAVideo,
|
||||
InferenceJobType::ImageGenApi => Self::ImageGenApi,
|
||||
InferenceJobType::Seedance2ProQueue => Self::Seedance2ProQueue,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -153,6 +158,7 @@ impl PublicInferenceJobType {
|
||||
Self::SoVitsSvc => InferenceJobType::SoVitsSvc,
|
||||
Self::StableDiffusion => InferenceJobType::StableDiffusion,
|
||||
Self::ImageGenApi => InferenceJobType::ImageGenApi,
|
||||
Self::Seedance2ProQueue => InferenceJobType::Seedance2ProQueue,
|
||||
Self::StyleTTS2 => InferenceJobType::StyleTTS2,
|
||||
Self::Tacotron2 => InferenceJobType::Tacotron2,
|
||||
Self::Unknown => InferenceJobType::Unknown,
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
[package]
|
||||
name = "seedance2-pro-job"
|
||||
edition = "2021"
|
||||
version = "0.0.1"
|
||||
authors = [
|
||||
"Brandon Thomas <bt@brand.io>",
|
||||
"Brandon Thomas <echelon@gmail.com>",
|
||||
]
|
||||
publish = false
|
||||
|
||||
[[bin]]
|
||||
name = "seedance2-pro-job"
|
||||
path = "src/main.rs"
|
||||
|
||||
[package.metadata.cargo-udeps.ignore]
|
||||
normal = ["workspace-build-acceleration"]
|
||||
|
||||
[dependencies]
|
||||
# Workspace hack (faster compile times)
|
||||
workspace-build-acceleration = { workspace = true }
|
||||
|
||||
# Internal
|
||||
actix_helpers.workspace = true
|
||||
bootstrap = { path = "../../../lib/bootstrap" }
|
||||
bucket_paths = { path = "../../../lib/bucket_paths" }
|
||||
cloud_storage.workspace = true
|
||||
concurrency = { path = "../../../lib/concurrency" }
|
||||
config = { path = "../../../lib/deprecated/config" }
|
||||
easyenv = { path = "../../../lib/easyenv" }
|
||||
enums.workspace = true
|
||||
errors.workspace = true
|
||||
hashing = { path = "../../../lib/files/hashing" }
|
||||
jobs_common = { path = "../../../lib/jobs_common" }
|
||||
mysql_queries = { path = "../../../schema/database/mysql_queries" }
|
||||
seedance2pro = { path = "../../../api_clients/seedance2pro" }
|
||||
server_environment = { path = "../../../lib/server_environment" }
|
||||
tokens.workspace = true
|
||||
|
||||
# External
|
||||
anyhow.workspace = true
|
||||
log.workspace = true
|
||||
|
||||
sqlx = { version = "0.7.4", features = [ "mysql", "runtime-tokio-rustls", "chrono" ] }
|
||||
|
||||
# Runtime
|
||||
tokio = { version = "1.40.0", features = ["macros", "sync"] }
|
||||
|
||||
# HTTP Client
|
||||
reqwest.workspace = true
|
||||
|
||||
# HTTP Server
|
||||
actix-web.workspace = true
|
||||
|
||||
# Serialization
|
||||
serde.workspace = true
|
||||
serde_derive.workspace = true
|
||||
serde_json.workspace = true
|
||||
@@ -0,0 +1,4 @@
|
||||
# default .env configurations (these can be overridden with k8s)
|
||||
|
||||
# Development
|
||||
SERVER_ENVIRONMENT='development'
|
||||
+101
@@ -0,0 +1,101 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use actix_web::error::ResponseError;
|
||||
use actix_web::http::header::ContentType;
|
||||
use actix_web::http::StatusCode;
|
||||
use actix_web::{web, HttpRequest, HttpResponse};
|
||||
use log::error;
|
||||
|
||||
use actix_helpers::response_serializers::error_to_json_http_response::error_to_json_http_response;
|
||||
|
||||
use crate::http_server::http_server_shared_state::HttpServerSharedState;
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct HealthCheckResponse {
|
||||
pub success: bool,
|
||||
pub is_healthy: bool,
|
||||
pub consecutive_failure_count: u64,
|
||||
pub consecutive_success_count: u64,
|
||||
pub total_failure_count: u64,
|
||||
pub total_success_count: u64,
|
||||
pub total_failure_ratio: f32,
|
||||
pub total_success_ratio: f32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub enum HealthCheckError {
|
||||
ServerError,
|
||||
}
|
||||
|
||||
impl ResponseError for HealthCheckError {
|
||||
fn status_code(&self) -> StatusCode {
|
||||
match *self {
|
||||
HealthCheckError::ServerError => StatusCode::INTERNAL_SERVER_ERROR,
|
||||
}
|
||||
}
|
||||
|
||||
fn error_response(&self) -> HttpResponse {
|
||||
error_to_json_http_response(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for HealthCheckError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{:?}", self)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_health_check_handler(
|
||||
_http_request: HttpRequest,
|
||||
server_state: web::Data<Arc<HttpServerSharedState>>,
|
||||
) -> Result<HttpResponse, HealthCheckError> {
|
||||
let job_stats = server_state
|
||||
.job_stats
|
||||
.get_status()
|
||||
.map_err(|e| {
|
||||
error!("Error serving health check status: {:?}", e);
|
||||
HealthCheckError::ServerError
|
||||
})?;
|
||||
|
||||
let total_tries = job_stats
|
||||
.total_failure_count
|
||||
.saturating_add(job_stats.total_success_count);
|
||||
|
||||
let total_success_ratio = if total_tries > 0 {
|
||||
(job_stats.total_success_count as f32) / (total_tries as f32)
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
let total_failure_ratio = if total_tries > 0 {
|
||||
1.0 - total_success_ratio
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
let is_healthy =
|
||||
job_stats.consecutive_failure_count < server_state.consecutive_failure_unhealthy_threshold;
|
||||
|
||||
let response = HealthCheckResponse {
|
||||
success: true,
|
||||
is_healthy,
|
||||
consecutive_failure_count: job_stats.consecutive_failure_count,
|
||||
consecutive_success_count: job_stats.consecutive_success_count,
|
||||
total_failure_count: job_stats.total_failure_count,
|
||||
total_success_count: job_stats.total_success_count,
|
||||
total_failure_ratio,
|
||||
total_success_ratio,
|
||||
};
|
||||
|
||||
let body = serde_json::to_string(&response).map_err(|_e| HealthCheckError::ServerError)?;
|
||||
|
||||
if is_healthy {
|
||||
Ok(HttpResponse::Ok()
|
||||
.content_type(ContentType::json())
|
||||
.body(body))
|
||||
} else {
|
||||
Ok(HttpResponse::InternalServerError()
|
||||
.content_type(ContentType::json())
|
||||
.body(body))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod health_check_handler;
|
||||
@@ -0,0 +1,7 @@
|
||||
use jobs_common::job_stats::JobStats;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct HttpServerSharedState {
|
||||
pub job_stats: JobStats,
|
||||
pub consecutive_failure_unhealthy_threshold: u64,
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
pub mod endpoints;
|
||||
pub mod http_server_shared_state;
|
||||
pub mod run_http_server;
|
||||
@@ -0,0 +1,67 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use actix_web::dev::Server;
|
||||
use actix_web::middleware::Logger;
|
||||
use actix_web::{web, App, HttpResponse, HttpServer};
|
||||
use log::info;
|
||||
|
||||
use bootstrap::bootstrap::ContainerEnvironment;
|
||||
use errors::AnyhowResult;
|
||||
use jobs_common::job_stats::JobStats;
|
||||
|
||||
use crate::http_server::endpoints::health_check_handler::get_health_check_handler;
|
||||
use crate::http_server::http_server_shared_state::HttpServerSharedState;
|
||||
|
||||
const DEFAULT_BIND_ADDRESS: &str = "0.0.0.0:12345";
|
||||
const DEFAULT_NUM_WORKERS: usize = 2;
|
||||
|
||||
pub struct CreateServerArgs {
|
||||
pub container_environment: ContainerEnvironment,
|
||||
pub job_stats: JobStats,
|
||||
}
|
||||
|
||||
pub fn run_http_server(args: CreateServerArgs) -> AnyhowResult<Server> {
|
||||
let bind_address = easyenv::get_env_string_or_default("HTTP_BIND_ADDRESS", DEFAULT_BIND_ADDRESS);
|
||||
let num_workers = easyenv::get_env_num("HTTP_NUM_WORKERS", DEFAULT_NUM_WORKERS)?;
|
||||
let _hostname = args.container_environment.hostname.clone();
|
||||
|
||||
let server_state = HttpServerSharedState {
|
||||
job_stats: args.job_stats.clone(),
|
||||
consecutive_failure_unhealthy_threshold: easyenv::get_env_num(
|
||||
"CONSECUTIVE_FAILURE_UNHEALTHY_THRESHOLD",
|
||||
3,
|
||||
)?,
|
||||
};
|
||||
|
||||
let server_state_arc = web::Data::new(Arc::new(server_state));
|
||||
|
||||
info!("Starting HTTP service (for k8s health checking).");
|
||||
|
||||
let log_format = "[%{HOSTNAME}e] IP=[%{X-Forwarded-For}i] \"%r\" %s %b \"%{Referer}i\" \"%{User-Agent}i\" %T";
|
||||
|
||||
let handle = HttpServer::new(move || {
|
||||
App::new()
|
||||
.app_data(server_state_arc.clone())
|
||||
.wrap(Logger::new(&log_format).exclude("/_status"))
|
||||
.service(
|
||||
web::resource("/")
|
||||
.route(web::get().to(|| HttpResponse::Ok()))
|
||||
.route(web::head().to(|| HttpResponse::Ok())),
|
||||
)
|
||||
.service(
|
||||
web::resource("/_status")
|
||||
.route(web::get().to(get_health_check_handler))
|
||||
.route(web::head().to(|| HttpResponse::Ok())),
|
||||
)
|
||||
})
|
||||
.bind(bind_address)?
|
||||
.workers(num_workers)
|
||||
.run();
|
||||
|
||||
Ok(handle)
|
||||
}
|
||||
|
||||
pub async fn launch_http_server(args: CreateServerArgs) -> AnyhowResult<()> {
|
||||
run_http_server(args)?.await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
|
||||
use log::{error, info, warn};
|
||||
use enums::by_table::generic_inference_jobs::frontend_failure_category::FrontendFailureCategory;
|
||||
use mysql_queries::queries::generic_inference::job::mark_job_failed_by_token::{mark_job_failed_by_token, MarkJobFailedByTokenArgs};
|
||||
use mysql_queries::queries::generic_inference::seedance2pro::list_pending_seedance2pro_jobs::list_pending_seedance2pro_jobs;
|
||||
use seedance2pro::requests::poll_orders::poll_orders::{poll_orders, PollOrdersArgs, TaskStatus};
|
||||
|
||||
use crate::job::job_loop::process_completed_order::{mark_job_failed, process_completed_order};
|
||||
use crate::job_dependencies::JobDependencies;
|
||||
|
||||
pub async fn main_loop(job_dependencies: JobDependencies) {
|
||||
while !job_dependencies.application_shutdown.get() {
|
||||
let result = run_poll_iteration(&job_dependencies).await;
|
||||
|
||||
if let Err(err) = result {
|
||||
error!("Error in poll iteration: {:?}", err);
|
||||
let _ = job_dependencies.job_stats.increment_failure_count();
|
||||
}
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(job_dependencies.poll_interval_millis)).await;
|
||||
}
|
||||
|
||||
warn!("Seedance2Pro job runner main loop is shut down.");
|
||||
}
|
||||
|
||||
async fn run_poll_iteration(deps: &JobDependencies) -> anyhow::Result<()> {
|
||||
// 1. Query all non-terminal Seedance2Pro jobs from DB.
|
||||
let pending_jobs = list_pending_seedance2pro_jobs(&deps.mysql_pool).await?;
|
||||
|
||||
if pending_jobs.is_empty() {
|
||||
info!("No pending Seedance2Pro jobs.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
info!("Found {} pending Seedance2Pro job(s).", pending_jobs.len());
|
||||
|
||||
// Build a lookup: order_id -> job
|
||||
let mut job_by_order_id: HashMap<String, _> = pending_jobs
|
||||
.into_iter()
|
||||
.map(|job| (job.order_id.clone(), job))
|
||||
.collect();
|
||||
|
||||
// 2. Poll all orders from seedance2pro API (exhausting pagination).
|
||||
let mut all_orders = Vec::new();
|
||||
let mut cursor: Option<u64> = None;
|
||||
|
||||
loop {
|
||||
let response = poll_orders(PollOrdersArgs {
|
||||
session: &deps.seedance2pro_session,
|
||||
cursor,
|
||||
})
|
||||
.await
|
||||
.map_err(|err| {
|
||||
warn!("Error polling seedance2pro orders: {:?}", err);
|
||||
anyhow::anyhow!("poll_orders failed: {:?}", err)
|
||||
})?;
|
||||
|
||||
info!("Polled {} orders (cursor={:?})", response.orders.len(), cursor);
|
||||
all_orders.extend(response.orders);
|
||||
|
||||
cursor = response.next_cursor;
|
||||
if cursor.is_none() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Match API orders to DB jobs and process terminal ones.
|
||||
for order in &all_orders {
|
||||
let job = match job_by_order_id.remove(&order.order_id) {
|
||||
Some(j) => j,
|
||||
None => continue, // Not one of our pending jobs.
|
||||
};
|
||||
|
||||
match &order.task_status {
|
||||
TaskStatus::Completed => {
|
||||
info!(
|
||||
"Order {} completed, processing job {}",
|
||||
order.order_id,
|
||||
job.job_token.as_str()
|
||||
);
|
||||
if let Err(err) = process_completed_order(deps, &job, order).await {
|
||||
warn!(
|
||||
"Error processing completed order {}: {:?}",
|
||||
order.order_id, err
|
||||
);
|
||||
let _ = deps.job_stats.increment_failure_count();
|
||||
} else {
|
||||
let _ = deps.job_stats.increment_success_count();
|
||||
}
|
||||
}
|
||||
TaskStatus::Failed => {
|
||||
let reason = order
|
||||
.fail_reason
|
||||
.as_deref()
|
||||
.unwrap_or("unknown failure reason");
|
||||
|
||||
let reason_lower = reason.to_lowercase();
|
||||
|
||||
let platform_rules_violation = reason_lower.contains("violates") ||
|
||||
reason_lower.contains("platform rules") ||
|
||||
reason_lower.contains("please modify");
|
||||
|
||||
let frontend_failure_category = if platform_rules_violation {
|
||||
Some(FrontendFailureCategory::ModelRulesViolation)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
warn!(
|
||||
"Order {} failed: {}. Marking job {} failed.",
|
||||
order.order_id, reason, job.job_token.as_str()
|
||||
);
|
||||
|
||||
let mark_failed_result = mark_job_failed_by_token(MarkJobFailedByTokenArgs {
|
||||
pool: &deps.mysql_pool,
|
||||
job_token: &job.job_token,
|
||||
maybe_public_failure_reason: Some(reason),
|
||||
internal_debugging_failure_reason: reason,
|
||||
maybe_frontend_failure_category: frontend_failure_category,
|
||||
}).await;
|
||||
|
||||
if let Err(err) = mark_failed_result {
|
||||
error!(
|
||||
"Error marking job {} as failed: {:?}",
|
||||
job.job_token.as_str(),
|
||||
err
|
||||
);
|
||||
}
|
||||
}
|
||||
TaskStatus::Pending | TaskStatus::Processing => {
|
||||
// Still in progress — check again next poll.
|
||||
}
|
||||
TaskStatus::Unknown(unknown_status) => {
|
||||
// NB: Keep polling it?
|
||||
warn!("Unknown order status: {:?}", unknown_status);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod main_loop;
|
||||
pub mod process_completed_order;
|
||||
@@ -0,0 +1,160 @@
|
||||
use anyhow::anyhow;
|
||||
use log::{error, info};
|
||||
|
||||
use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
|
||||
use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType;
|
||||
use enums::by_table::media_files::media_file_class::MediaFileClass;
|
||||
use enums::by_table::media_files::media_file_origin_category::MediaFileOriginCategory;
|
||||
use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory;
|
||||
use enums::by_table::media_files::media_file_type::MediaFileType;
|
||||
use errors::AnyhowResult;
|
||||
use hashing::sha256::sha256_hash_bytes::sha256_hash_bytes;
|
||||
use mysql_queries::queries::generic_inference::job::mark_job_failed_by_token::{mark_job_failed_by_token, MarkJobFailedByTokenArgs};
|
||||
use mysql_queries::queries::generic_inference::seedance2pro::list_pending_seedance2pro_jobs::PendingSeedance2ProJob;
|
||||
use mysql_queries::queries::media_files::create::insert_builder::media_file_insert_builder::MediaFileInsertBuilder;
|
||||
use mysql_queries::queries::generic_inference::web::mark_generic_inference_job_successfully_done_by_token::mark_generic_inference_job_successfully_done_by_token;
|
||||
use seedance2pro::requests::poll_orders::poll_orders::OrderStatus;
|
||||
use crate::job_dependencies::JobDependencies;
|
||||
|
||||
const PREFIX : &str = "artcraft_";
|
||||
const SUFFIX : &str = ".mp4";
|
||||
|
||||
/// Download the completed video, upload to bucket, create media file record, and mark job done.
|
||||
pub async fn process_completed_order(
|
||||
deps: &JobDependencies,
|
||||
job: &PendingSeedance2ProJob,
|
||||
order: &OrderStatus,
|
||||
) -> AnyhowResult<()> {
|
||||
// Get the video URL.
|
||||
let video_url = match &order.result_url {
|
||||
Some(url) => url.as_str(),
|
||||
None => {
|
||||
// Fall back to the first result entry if the top-level result_url is missing.
|
||||
match order.results.first() {
|
||||
Some(result) => result.url.as_str(),
|
||||
None => {
|
||||
return Err(anyhow!(
|
||||
"Completed order {} has no result_url and no results",
|
||||
order.order_id
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
info!(
|
||||
"Downloading video for order {} from: {}",
|
||||
order.order_id, video_url
|
||||
);
|
||||
|
||||
// Download the video bytes.
|
||||
let video_bytes: Vec<u8> = reqwest::get(video_url)
|
||||
.await
|
||||
.map_err(|err| anyhow!("reqwest error downloading video: {:?}", err))?
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|err| anyhow!("error reading video bytes: {:?}", err))?
|
||||
.to_vec();
|
||||
|
||||
info!(
|
||||
"Downloaded {} bytes for order {}",
|
||||
video_bytes.len(),
|
||||
order.order_id
|
||||
);
|
||||
|
||||
// Hash the video.
|
||||
let checksum = sha256_hash_bytes(&video_bytes)
|
||||
.map_err(|err| anyhow!("error hashing video: {:?}", err))?;
|
||||
|
||||
// Build the bucket path.
|
||||
let bucket_path = MediaFileBucketPath::generate_new(Some(PREFIX), Some(SUFFIX));
|
||||
|
||||
let object_path = bucket_path.get_full_object_path_str();
|
||||
|
||||
info!(
|
||||
"Uploading video to public bucket at path: {}",
|
||||
object_path
|
||||
);
|
||||
|
||||
// Upload to public bucket.
|
||||
deps
|
||||
.public_bucket_client
|
||||
.upload_file_with_content_type_process(object_path, &video_bytes, "video/mp4")
|
||||
.await
|
||||
.map_err(|err| anyhow!("error uploading video to bucket: {:?}", err))?;
|
||||
|
||||
info!(
|
||||
"Uploaded video for order {}. Creating media file record.",
|
||||
order.order_id
|
||||
);
|
||||
|
||||
// Optionally extract frame dimensions from the order results.
|
||||
let maybe_frame_width = order.results.first().map(|r| r.width);
|
||||
let maybe_frame_height = order.results.first().map(|r| r.height);
|
||||
|
||||
// Insert media file record.
|
||||
let media_file_token = MediaFileInsertBuilder::new()
|
||||
.maybe_creator_user(job.maybe_creator_user_token.as_ref())
|
||||
.maybe_creator_anonymous_visitor(job.maybe_creator_anonymous_visitor_token.as_ref())
|
||||
.creator_ip_address(&job.creator_ip_address)
|
||||
.creator_set_visibility(job.creator_set_visibility)
|
||||
.media_file_class(MediaFileClass::Video)
|
||||
.media_file_type(MediaFileType::Mp4)
|
||||
.media_file_origin_category(MediaFileOriginCategory::Inference)
|
||||
.media_file_origin_product_category(MediaFileOriginProductCategory::VideoGeneration)
|
||||
.mime_type("video/mp4")
|
||||
.file_size_bytes(video_bytes.len() as u64)
|
||||
.maybe_frame_width(maybe_frame_width)
|
||||
.maybe_frame_height(maybe_frame_height)
|
||||
.checksum_sha2(&checksum)
|
||||
.maybe_prompt_token(job.maybe_prompt_token.as_ref())
|
||||
.public_bucket_directory_hash(&bucket_path)
|
||||
.insert_pool(&deps.mysql_pool)
|
||||
.await
|
||||
.map_err(|err| anyhow!("error inserting media file record: {:?}", err))?;
|
||||
|
||||
info!(
|
||||
"Created media file {} for order {}. Marking job {} complete.",
|
||||
media_file_token.as_str(),
|
||||
order.order_id,
|
||||
job.job_token.as_str()
|
||||
);
|
||||
|
||||
// Mark inference job as successfully completed.
|
||||
mark_generic_inference_job_successfully_done_by_token(
|
||||
&deps.mysql_pool,
|
||||
&job.job_token,
|
||||
Some(InferenceResultType::MediaFile),
|
||||
Some(media_file_token.as_str()),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
error!(
|
||||
"Error marking job {} done: {:?}",
|
||||
job.job_token.as_str(),
|
||||
err
|
||||
);
|
||||
anyhow!("error marking job done: {:?}", err)
|
||||
})?;
|
||||
|
||||
info!("Job {} completed successfully.", job.job_token.as_str());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Mark a job as permanently failed.
|
||||
pub async fn mark_job_failed(
|
||||
deps: &JobDependencies,
|
||||
job: &PendingSeedance2ProJob,
|
||||
fail_reason: &str,
|
||||
) -> AnyhowResult<()> {
|
||||
mark_job_failed_by_token(MarkJobFailedByTokenArgs {
|
||||
pool: &deps.mysql_pool,
|
||||
job_token: &job.job_token,
|
||||
maybe_public_failure_reason: Some(fail_reason),
|
||||
internal_debugging_failure_reason: fail_reason,
|
||||
maybe_frontend_failure_category: None,
|
||||
}).await
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod job_loop;
|
||||
@@ -0,0 +1,26 @@
|
||||
use cloud_storage::bucket_client::BucketClient;
|
||||
use concurrency::relaxed_atomic_bool::RelaxedAtomicBool;
|
||||
use jobs_common::job_stats::JobStats;
|
||||
use seedance2pro::creds::seedance2pro_session::Seedance2ProSession;
|
||||
use server_environment::ServerEnvironment;
|
||||
use sqlx::MySqlPool;
|
||||
|
||||
pub struct JobDependencies {
|
||||
pub mysql_pool: MySqlPool,
|
||||
|
||||
/// Public GCS/S3 bucket for storing generated videos.
|
||||
pub public_bucket_client: BucketClient,
|
||||
|
||||
/// Session credentials for polling seedance2-pro.com.
|
||||
pub seedance2pro_session: Seedance2ProSession,
|
||||
|
||||
pub server_environment: ServerEnvironment,
|
||||
|
||||
pub job_stats: JobStats,
|
||||
|
||||
/// How long to sleep between poll iterations (milliseconds).
|
||||
pub poll_interval_millis: u64,
|
||||
|
||||
/// Set to `true` from another thread to trigger graceful shutdown.
|
||||
pub application_shutdown: RelaxedAtomicBool,
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
// Never allow these
|
||||
#![forbid(private_bounds)]
|
||||
#![forbid(private_interfaces)]
|
||||
#![forbid(unused_must_use)]
|
||||
|
||||
// Always allow
|
||||
#![allow(dead_code)]
|
||||
#![allow(non_snake_case)]
|
||||
|
||||
#[macro_use] extern crate serde_derive;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::anyhow;
|
||||
use log::{info, warn};
|
||||
use sqlx::mysql::MySqlPoolOptions;
|
||||
|
||||
use bootstrap::bootstrap::{bootstrap, BootstrapArgs};
|
||||
use cloud_storage::bucket_client::BucketClient;
|
||||
use concurrency::relaxed_atomic_bool::RelaxedAtomicBool;
|
||||
use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING;
|
||||
use config::shared_constants::DEFAULT_RUST_LOG;
|
||||
use errors::AnyhowResult;
|
||||
use jobs_common::job_stats::JobStats;
|
||||
use seedance2pro::creds::seedance2pro_session::Seedance2ProSession;
|
||||
use server_environment::ServerEnvironment;
|
||||
|
||||
use crate::http_server::run_http_server::{launch_http_server, CreateServerArgs};
|
||||
use crate::job::job_loop::main_loop::main_loop;
|
||||
use crate::job_dependencies::JobDependencies;
|
||||
|
||||
pub mod http_server;
|
||||
pub mod job;
|
||||
pub mod job_dependencies;
|
||||
|
||||
// Bucket config
|
||||
const ENV_ACCESS_KEY: &str = "ACCESS_KEY";
|
||||
const ENV_SECRET_KEY: &str = "SECRET_KEY";
|
||||
const ENV_REGION_NAME: &str = "REGION_NAME";
|
||||
const ENV_PUBLIC_BUCKET_NAME: &str = "PUBLIC_BUCKET_NAME";
|
||||
const ENV_S3_ENDPOINT: &str = "S3_COMPATIBLE_ENDPOINT_URL";
|
||||
const ENV_SEEDANCE2PRO_COOKIES : &str = "SEEDANCE2PRO_COOKIES";
|
||||
|
||||
#[actix_web::main]
|
||||
async fn main() -> AnyhowResult<()> {
|
||||
|
||||
let container_environment = bootstrap(BootstrapArgs {
|
||||
app_name: "seedance2-pro-job",
|
||||
default_logging_override: Some(DEFAULT_RUST_LOG),
|
||||
config_search_directories: &[".", "./config", "crates/service/job/seedance2_pro_job/config"],
|
||||
ignore_legacy_dot_env_file: true,
|
||||
})?;
|
||||
|
||||
info!("Hostname: {}", &container_environment.hostname);
|
||||
|
||||
let _k8s_node_name = easyenv::get_env_string_optional("K8S_NODE_NAME");
|
||||
let _k8s_pod_name = easyenv::get_env_string_optional("K8S_POD_NAME");
|
||||
|
||||
let db_connection_string = easyenv::get_env_string_or_default(
|
||||
"MYSQL_URL",
|
||||
DEFAULT_MYSQL_CONNECTION_STRING,
|
||||
);
|
||||
|
||||
info!("Connecting to database...");
|
||||
|
||||
let mysql_pool = MySqlPoolOptions::new()
|
||||
.max_connections(2)
|
||||
.connect(&db_connection_string)
|
||||
.await?;
|
||||
|
||||
info!("Connected to MySQL.");
|
||||
|
||||
let server_environment = ServerEnvironment::from_str(
|
||||
&easyenv::get_env_string_required("SERVER_ENVIRONMENT")?,
|
||||
)
|
||||
.ok_or(anyhow!("invalid server environment"))?;
|
||||
|
||||
// Bucket setup
|
||||
let access_key = easyenv::get_env_string_required(ENV_ACCESS_KEY)?;
|
||||
let secret_key = easyenv::get_env_string_required(ENV_SECRET_KEY)?;
|
||||
let region_name = easyenv::get_env_string_required(ENV_REGION_NAME)?;
|
||||
let public_bucket_name = easyenv::get_env_string_required(ENV_PUBLIC_BUCKET_NAME)?;
|
||||
let s3_compatible_endpoint_url = easyenv::get_env_string_required(ENV_S3_ENDPOINT)?;
|
||||
|
||||
let bucket_timeout = easyenv::get_env_duration_seconds_or_default(
|
||||
"BUCKET_TIMEOUT_SECONDS",
|
||||
Duration::from_secs(60 * 5),
|
||||
);
|
||||
|
||||
let public_bucket_client = BucketClient::create(
|
||||
&access_key,
|
||||
&secret_key,
|
||||
®ion_name,
|
||||
&public_bucket_name,
|
||||
&s3_compatible_endpoint_url,
|
||||
None,
|
||||
Some(bucket_timeout),
|
||||
)?;
|
||||
|
||||
// Seedance2Pro session from cookie string
|
||||
let seedance2pro_cookies = easyenv::get_env_string_required(ENV_SEEDANCE2PRO_COOKIES)?;
|
||||
let seedance2pro_session = Seedance2ProSession::from_cookies_string(seedance2pro_cookies);
|
||||
|
||||
// How often to poll for results (default: 15 seconds)
|
||||
let poll_interval_millis: u64 = easyenv::get_env_num(
|
||||
"SEEDANCE_POLL_INTERVAL_MILLIS",
|
||||
5_000,
|
||||
)?;
|
||||
|
||||
let application_shutdown = RelaxedAtomicBool::new(false);
|
||||
let job_stats = JobStats::new();
|
||||
|
||||
let create_server_args = CreateServerArgs {
|
||||
container_environment: container_environment.clone(),
|
||||
job_stats: job_stats.clone(),
|
||||
};
|
||||
|
||||
let job_dependencies = JobDependencies {
|
||||
mysql_pool,
|
||||
public_bucket_client,
|
||||
seedance2pro_session,
|
||||
server_environment,
|
||||
job_stats,
|
||||
poll_interval_millis,
|
||||
application_shutdown: application_shutdown.clone(),
|
||||
};
|
||||
|
||||
std::thread::spawn(move || {
|
||||
let actix_runtime = actix_web::rt::System::new();
|
||||
let http_server_handle = launch_http_server(create_server_args);
|
||||
|
||||
actix_runtime.block_on(http_server_handle)
|
||||
.expect("HTTP server should not exit.");
|
||||
|
||||
warn!("Server thread is shut down.");
|
||||
application_shutdown.set(true);
|
||||
});
|
||||
|
||||
main_loop(job_dependencies).await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -52,6 +52,7 @@ email_sender= { path = "../../../lib/email_sender" }
|
||||
enums.workspace = true
|
||||
enums_public.workspace = true
|
||||
errors.workspace = true
|
||||
seedance2pro.workspace = true
|
||||
fal_client.workspace = true
|
||||
filesys.workspace = true
|
||||
google_sign_in = { path = "../../../api_clients/google_sign_in" }
|
||||
@@ -99,7 +100,7 @@ once_cell.workspace = true
|
||||
rand.workspace = true
|
||||
regex.workspace = true
|
||||
strum_macros = "0.26.4"
|
||||
time = "0.3.30"
|
||||
time = "0.3.36"
|
||||
url.workspace = true
|
||||
|
||||
# Money
|
||||
@@ -114,7 +115,11 @@ language-tags = "0.3.2" # Parsing single tags richly
|
||||
reqwest.workspace = true
|
||||
|
||||
# Email
|
||||
resend-rs = "0.21.0"
|
||||
resend-rs = { version = "0.21.0", default-features = false, features = ["rustls-tls"] }
|
||||
# NB: Both ring (via reqwest 0.12/hyper-rustls) and aws-lc-rs (via quinn/resend-rs)
|
||||
# are features of rustls 0.23 in this binary. rustls requires an explicit default
|
||||
# provider when multiple crypto providers are compiled in.
|
||||
rustls = { version = "0.23", features = ["ring"], default-features = false }
|
||||
|
||||
# Runtime
|
||||
tokio.workspace = true
|
||||
|
||||
@@ -42,7 +42,8 @@ GC_ENABLED_PUBLIC_BUCKET_NAME="REPLACE_ME"
|
||||
AUDIO_UPLOADS_BUCKET_ROOT=/user_uploaded_w2l_audio
|
||||
|
||||
# Fake API key for FAL
|
||||
FAL_WEBHOOK_URL="https://macaroni-1.tailce84f.ts.net/v1/webhooks/fal"
|
||||
FAL_WEBHOOK_URL="https://brandons-macbook-pro.taild62114.ts.net/v1/webhooks/fal"
|
||||
#FAL_WEBHOOK_URL="https://macaroni-1.tailce84f.ts.net/v1/webhooks/fal"
|
||||
#FAL_WEBHOOK_URL="https://halide.tailce84f.ts.net/v1/webhooks/fal"
|
||||
#FAL_WEBHOOK_URL="https://s1.pancake-chimaera.ts.net/api/fal/webhook"
|
||||
|
||||
|
||||
+1
@@ -1,5 +1,6 @@
|
||||
pub mod kling_2p5_turbo_pro_multi_function_video_gen_handler;
|
||||
pub mod kling_2p6_pro_multi_function_video_gen_handler;
|
||||
pub mod seedance_2p0_multi_function_video_gen_handler;
|
||||
pub mod sora_2_multi_function_video_gen_handler;
|
||||
pub mod sora_2_pro_multi_function_video_gen_handler;
|
||||
pub mod veo_3p1_fast_multi_function_video_gen_handler;
|
||||
|
||||
+390
@@ -0,0 +1,390 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error;
|
||||
use crate::http_server::common_responses::common_web_error::CommonWebError;
|
||||
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 crate::util::http_download_url_to_bytes::http_download_url_to_bytes;
|
||||
use crate::util::lookup::lookup_image_urls_as_map::lookup_image_urls_as_map;
|
||||
use actix_web::web::Json;
|
||||
use actix_web::{web, HttpRequest};
|
||||
use artcraft_api_defs::generate::video::multi_function::seedance_2p0_multi_function_video_gen::{
|
||||
Seedance2p0AspectRatio, Seedance2p0BatchCount, Seedance2p0MultiFunctionVideoGenRequest,
|
||||
Seedance2p0MultiFunctionVideoGenResponse,
|
||||
};
|
||||
use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptContextSemanticType;
|
||||
use enums::by_table::prompts::prompt_type::PromptType;
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::common::model_type::ModelType;
|
||||
use enums::common::visibility::Visibility;
|
||||
use http_server_common::request::get_request_ip::get_request_ip;
|
||||
use log::{error, info, warn};
|
||||
use mysql_queries::queries::generic_inference::seedance2pro::insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token::{
|
||||
insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token,
|
||||
InsertGenericInferenceForSeedance2ProWithAprioriJobTokenArgs,
|
||||
};
|
||||
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 seedance2pro::creds::seedance2pro_session::Seedance2ProSession;
|
||||
use seedance2pro::requests::generate_video::generate_video::{
|
||||
generate_video, BatchCount, GenerateVideoArgs, Resolution,
|
||||
};
|
||||
use seedance2pro::requests::prepare_image_upload::prepare_image_upload::{
|
||||
prepare_image_upload, PrepareImageUploadArgs,
|
||||
};
|
||||
use seedance2pro::requests::upload_image::upload_image::{upload_image, UploadImageArgs};
|
||||
use sqlx::Acquire;
|
||||
use tokens::tokens::generic_inference_jobs::InferenceJobToken;
|
||||
|
||||
/// Seedance 2.0 Multi-Function video generation (text-to-video, keyframe, and reference).
|
||||
#[utoipa::path(
|
||||
post,
|
||||
tag = "Generate Video (Multi-Function)",
|
||||
path = "/v1/generate/video/multi_function/seedance_2p0",
|
||||
responses(
|
||||
(status = 200, description = "Success", body = Seedance2p0MultiFunctionVideoGenResponse),
|
||||
),
|
||||
params(
|
||||
("request" = Seedance2p0MultiFunctionVideoGenRequest, description = "Payload for Request"),
|
||||
)
|
||||
)]
|
||||
pub async fn seedance_2p0_multi_function_video_gen_handler(
|
||||
http_request: HttpRequest,
|
||||
request: Json<Seedance2p0MultiFunctionVideoGenRequest>,
|
||||
server_state: web::Data<Arc<ServerState>>,
|
||||
) -> Result<Json<Seedance2p0MultiFunctionVideoGenResponse>, CommonWebError> {
|
||||
|
||||
payments_error_test(&request.prompt.as_deref().unwrap_or(""))?;
|
||||
|
||||
let mut mysql_connection = server_state.mysql_pool
|
||||
.acquire()
|
||||
.await?;
|
||||
|
||||
let maybe_user_session = server_state
|
||||
.session_checker
|
||||
.maybe_get_user_session_from_connection(&http_request, &mut mysql_connection)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("Session checker error: {:?}", e);
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
let maybe_avt_token = server_state
|
||||
.avt_cookie_manager
|
||||
.get_avt_token_from_request(&http_request);
|
||||
|
||||
let user_token = match maybe_user_session.as_ref() {
|
||||
Some(session) => &session.user_token,
|
||||
None => {
|
||||
return Err(CommonWebError::NotAuthorized);
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) {
|
||||
return Err(CommonWebError::BadInputWithSimpleMessage(reason));
|
||||
}
|
||||
|
||||
// --- Collect all media tokens to look up ---
|
||||
|
||||
let mut all_media_tokens = Vec::new();
|
||||
if let Some(token) = request.start_frame_media_token.as_ref() {
|
||||
all_media_tokens.push(token.clone());
|
||||
}
|
||||
if let Some(token) = request.end_frame_media_token.as_ref() {
|
||||
all_media_tokens.push(token.clone());
|
||||
}
|
||||
if let Some(tokens) = request.reference_image_media_tokens.as_ref() {
|
||||
all_media_tokens.extend(tokens.iter().cloned());
|
||||
}
|
||||
|
||||
let image_urls_by_token = if all_media_tokens.is_empty() {
|
||||
std::collections::HashMap::new()
|
||||
} else {
|
||||
info!("Looking up image media tokens: {:?}", all_media_tokens);
|
||||
lookup_image_urls_as_map(
|
||||
&http_request,
|
||||
&mut mysql_connection,
|
||||
server_state.server_environment,
|
||||
&all_media_tokens,
|
||||
).await?
|
||||
};
|
||||
|
||||
// --- Insert idempotency token ---
|
||||
|
||||
insert_idempotency_token(&request.uuid_idempotency_token, &mut *mysql_connection)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
error!("Error inserting idempotency token: {:?}", err);
|
||||
CommonWebError::BadInputWithSimpleMessage("repeated idempotency token".to_string())
|
||||
})?;
|
||||
|
||||
// --- Build seedance2pro session ---
|
||||
|
||||
let session = Seedance2ProSession::from_cookies_string(
|
||||
server_state.seedance2pro.cookies.clone()
|
||||
);
|
||||
|
||||
// --- Upload images to seedance2pro CDN ---
|
||||
|
||||
let start_frame_url = match request.start_frame_media_token.as_ref() {
|
||||
None => None,
|
||||
Some(token) => match image_urls_by_token.get(token) {
|
||||
None => return Err(CommonWebError::BadInputWithSimpleMessage("Start frame media not found.".to_string())),
|
||||
Some(url) => Some(upload_to_seedance2pro(&session, url).await?),
|
||||
}
|
||||
};
|
||||
|
||||
let end_frame_url = match request.end_frame_media_token.as_ref() {
|
||||
None => None,
|
||||
Some(token) => match image_urls_by_token.get(token) {
|
||||
None => return Err(CommonWebError::BadInputWithSimpleMessage("End frame media not found.".to_string())),
|
||||
Some(url) => Some(upload_to_seedance2pro(&session, url).await?),
|
||||
}
|
||||
};
|
||||
|
||||
let reference_image_urls: Option<Vec<String>> = match request.reference_image_media_tokens.as_ref() {
|
||||
None => None,
|
||||
Some(tokens) if tokens.is_empty() => None,
|
||||
Some(tokens) => {
|
||||
let mut urls = Vec::with_capacity(tokens.len());
|
||||
for token in tokens {
|
||||
match image_urls_by_token.get(token) {
|
||||
None => return Err(CommonWebError::BadInputWithSimpleMessage("Reference image media not found.".to_string())),
|
||||
Some(url) => {
|
||||
let seedance_url = upload_to_seedance2pro(&session, url).await?;
|
||||
urls.push(seedance_url);
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(urls)
|
||||
}
|
||||
};
|
||||
|
||||
// --- Map request params to seedance2pro types ---
|
||||
|
||||
let resolution = match request.aspect_ratio {
|
||||
Some(Seedance2p0AspectRatio::Landscape16x9) => Resolution::Landscape16x9,
|
||||
Some(Seedance2p0AspectRatio::Portrait9x16) => Resolution::Portrait9x16,
|
||||
Some(Seedance2p0AspectRatio::Square1x1) => Resolution::Square1x1,
|
||||
Some(Seedance2p0AspectRatio::Standard4x3) => Resolution::Standard4x3,
|
||||
Some(Seedance2p0AspectRatio::Portrait3x4) => Resolution::Portrait3x4,
|
||||
None => Resolution::Landscape16x9,
|
||||
};
|
||||
|
||||
let batch_count = match request.batch_count {
|
||||
Some(Seedance2p0BatchCount::One) | None => BatchCount::One,
|
||||
Some(Seedance2p0BatchCount::Two) => BatchCount::Two,
|
||||
Some(Seedance2p0BatchCount::Four) => BatchCount::Four,
|
||||
};
|
||||
|
||||
let duration_seconds = request.duration_seconds.unwrap_or(5).clamp(4, 15);
|
||||
let prompt = request.prompt.clone().unwrap_or_else(|| "".to_string());
|
||||
|
||||
let video_gen_args = GenerateVideoArgs {
|
||||
session: &session,
|
||||
prompt,
|
||||
resolution,
|
||||
duration_seconds,
|
||||
batch_count,
|
||||
start_frame_url,
|
||||
end_frame_url,
|
||||
reference_image_urls,
|
||||
};
|
||||
|
||||
// --- Calculate cost and charge wallet ---
|
||||
|
||||
let cost_in_cents = video_gen_args.estimate_cost_in_usd_cents(); // NB: ArtCraft credits are 1:1 with USD cents.
|
||||
|
||||
let apriori_job_token = InferenceJobToken::generate();
|
||||
|
||||
info!("Charging wallet: {} cents ({} credits)", cost_in_cents, cost_in_cents);
|
||||
|
||||
attempt_wallet_deduction_else_common_web_error(
|
||||
user_token,
|
||||
Some(apriori_job_token.as_str()),
|
||||
cost_in_cents,
|
||||
&mut mysql_connection,
|
||||
).await?;
|
||||
|
||||
// --- Call seedance2pro API ---
|
||||
|
||||
let gen_response = generate_video(video_gen_args)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
warn!("Error calling seedance2pro generate_video: {:?}", err);
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
info!(
|
||||
"Seedance2pro task_id={}, order_id={}",
|
||||
gen_response.task_id, gen_response.order_id
|
||||
);
|
||||
|
||||
// --- DB writes in a 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);
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
// NB: Don't fail the job if the prompt insert fails.
|
||||
let prompt_result = insert_prompt(InsertPromptArgs {
|
||||
maybe_apriori_prompt_token: None,
|
||||
prompt_type: PromptType::ArtcraftApp,
|
||||
maybe_creator_user_token: maybe_user_session.as_ref().map(|s| &s.user_token),
|
||||
maybe_model_type: Some(ModelType::Seedance2p0),
|
||||
maybe_generation_provider: Some(GenerationProvider::Artcraft),
|
||||
maybe_positive_prompt: request.prompt.as_deref(),
|
||||
maybe_negative_prompt: None,
|
||||
maybe_other_args: None,
|
||||
creator_ip_address: &ip_address,
|
||||
mysql_executor: &mut *transaction,
|
||||
phantom: Default::default(),
|
||||
}).await;
|
||||
|
||||
let prompt_token = match prompt_result {
|
||||
Ok(token) => Some(token),
|
||||
Err(err) => {
|
||||
warn!("Error inserting prompt: {:?}", err);
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(token) = prompt_token.as_ref() {
|
||||
let mut context_items: Vec<PromptContextItem> = Vec::new();
|
||||
|
||||
if let Some(media_token) = &request.start_frame_media_token {
|
||||
context_items.push(PromptContextItem {
|
||||
media_token: media_token.clone(),
|
||||
context_semantic_type: PromptContextSemanticType::VidStartFrame,
|
||||
});
|
||||
}
|
||||
if let Some(media_token) = &request.end_frame_media_token {
|
||||
context_items.push(PromptContextItem {
|
||||
media_token: media_token.clone(),
|
||||
context_semantic_type: PromptContextSemanticType::VidEndFrame,
|
||||
});
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Insert one DB job per order_id in the batch.
|
||||
let order_ids: Vec<String> = match gen_response.order_ids {
|
||||
Some(ids) if !ids.is_empty() => ids,
|
||||
_ => vec![gen_response.order_id],
|
||||
};
|
||||
|
||||
let mut all_job_tokens: Vec<InferenceJobToken> = Vec::with_capacity(order_ids.len());
|
||||
|
||||
for (i, order_id) in order_ids.iter().enumerate() {
|
||||
let job_token = if i == 0 {
|
||||
apriori_job_token.clone()
|
||||
} else {
|
||||
InferenceJobToken::generate()
|
||||
};
|
||||
|
||||
// For batch jobs beyond the first, generate a unique idempotency key.
|
||||
let idempotency_str = if i == 0 {
|
||||
request.uuid_idempotency_token.clone()
|
||||
} else {
|
||||
format!("{}-batch-{}", request.uuid_idempotency_token, i)
|
||||
};
|
||||
|
||||
let db_result = insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token(
|
||||
InsertGenericInferenceForSeedance2ProWithAprioriJobTokenArgs {
|
||||
apriori_job_token: &job_token,
|
||||
uuid_idempotency_token: &idempotency_str,
|
||||
maybe_external_third_party_id: order_id,
|
||||
maybe_inference_args: None,
|
||||
maybe_prompt_token: prompt_token.as_ref(),
|
||||
maybe_creator_user_token: maybe_user_session.as_ref().map(|s| &s.user_token),
|
||||
maybe_avt_token: maybe_avt_token.as_ref(),
|
||||
creator_ip_address: &ip_address,
|
||||
creator_set_visibility: Visibility::Public,
|
||||
mysql_executor: &mut *transaction,
|
||||
phantom: Default::default(),
|
||||
}
|
||||
).await;
|
||||
|
||||
match db_result {
|
||||
Ok(token) => {
|
||||
all_job_tokens.push(token);
|
||||
}
|
||||
Err(err) => {
|
||||
warn!("Error inserting seedance2pro inference job (order_id={}): {:?}", order_id, err);
|
||||
if i == 0 {
|
||||
return Err(CommonWebError::ServerError);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let first_job_token = all_job_tokens.first().cloned().ok_or_else(|| {
|
||||
error!("No inference job token was created");
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|err| {
|
||||
error!("Error committing MySQL transaction: {:?}", err);
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
Ok(Json(Seedance2p0MultiFunctionVideoGenResponse {
|
||||
success: true,
|
||||
inference_job_token: first_job_token,
|
||||
all_inference_job_tokens: all_job_tokens,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn upload_to_seedance2pro(
|
||||
session: &Seedance2ProSession,
|
||||
our_cdn_url: &str,
|
||||
) -> Result<String, CommonWebError> {
|
||||
let image_bytes = http_download_url_to_bytes(our_cdn_url)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
warn!("Error downloading media file from CDN: {:?}", err);
|
||||
CommonWebError::ServerError
|
||||
})?
|
||||
.to_vec();
|
||||
|
||||
let prepare_result = prepare_image_upload(PrepareImageUploadArgs { session })
|
||||
.await
|
||||
.map_err(|err| {
|
||||
warn!("Error preparing seedance2pro image upload: {:?}", err);
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
let upload_result = upload_image(UploadImageArgs {
|
||||
upload_url: prepare_result.upload_url,
|
||||
image_bytes,
|
||||
})
|
||||
.await
|
||||
.map_err(|err| {
|
||||
warn!("Error uploading image to seedance2pro: {:?}", err);
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
Ok(upload_result.public_url)
|
||||
}
|
||||
@@ -22,6 +22,7 @@ pub mod trending;
|
||||
pub mod tts;
|
||||
pub mod user_bookmarks;
|
||||
pub mod user_ratings;
|
||||
pub mod user_wallets;
|
||||
pub mod users;
|
||||
pub mod voice_clone_requests;
|
||||
pub mod voice_conversion;
|
||||
|
||||
+103
@@ -0,0 +1,103 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::http_server::common_responses::common_web_error::CommonWebError;
|
||||
use crate::state::server_state::ServerState;
|
||||
use actix_web::web::Json;
|
||||
use actix_web::{web, HttpRequest};
|
||||
use enums::common::payments_namespace::PaymentsNamespace;
|
||||
use log::{error, warn};
|
||||
use mysql_queries::queries::wallets::create_new_artcraft_wallet_for_owner_user::create_new_artcraft_wallet_for_owner_user;
|
||||
use mysql_queries::queries::wallets::find_primary_wallet_token_for_owner::find_primary_wallet_token_for_owner_using_connection;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::Acquire;
|
||||
use tokens::tokens::wallets::WalletToken;
|
||||
use utoipa::ToSchema;
|
||||
|
||||
#[derive(Serialize, Deserialize, ToSchema)]
|
||||
pub struct GetOrCreateSessionArtcraftWalletResponse {
|
||||
pub success: bool,
|
||||
pub wallet_token: WalletToken,
|
||||
}
|
||||
|
||||
/// NB: This endpoint was created primarily for local testing with new dev accounts
|
||||
/// It should be harmless to have in production as wallet creation is idempotent
|
||||
/// and (should be) side effect free.
|
||||
#[utoipa::path(
|
||||
get,
|
||||
tag = "Wallets",
|
||||
path = "/v1/wallets/session_artcraft_wallet",
|
||||
responses(
|
||||
(status = 200, description = "Success", body = GetOrCreateSessionArtcraftWalletResponse),
|
||||
),
|
||||
)]
|
||||
pub async fn get_or_create_session_artcraft_wallet_handler(
|
||||
http_request: HttpRequest,
|
||||
server_state: web::Data<Arc<ServerState>>,
|
||||
) -> Result<Json<GetOrCreateSessionArtcraftWalletResponse>, CommonWebError> {
|
||||
let mut mysql_connection = server_state.mysql_pool
|
||||
.acquire()
|
||||
.await
|
||||
.map_err(|err| {
|
||||
error!("Error acquiring MySQL connection: {:?}", err);
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
let maybe_user_session = server_state
|
||||
.session_checker
|
||||
.maybe_get_user_session_from_connection(&http_request, &mut mysql_connection)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("Session checker error: {:?}", e);
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
let user_token = match maybe_user_session {
|
||||
Some(session) => session.user_token,
|
||||
None => return Err(CommonWebError::NotAuthorized),
|
||||
};
|
||||
|
||||
let maybe_wallet_token = find_primary_wallet_token_for_owner_using_connection(
|
||||
&user_token,
|
||||
PaymentsNamespace::Artcraft,
|
||||
&mut mysql_connection,
|
||||
).await.map_err(|err| {
|
||||
error!("Error finding primary artcraft wallet for user {:?}: {:?}", user_token, err);
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
if let Some(wallet_token) = maybe_wallet_token {
|
||||
return Ok(Json(GetOrCreateSessionArtcraftWalletResponse {
|
||||
success: true,
|
||||
wallet_token,
|
||||
}));
|
||||
}
|
||||
|
||||
let mut transaction = mysql_connection
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|err| {
|
||||
error!("Error starting MySQL transaction: {:?}", err);
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
let wallet_token = create_new_artcraft_wallet_for_owner_user(
|
||||
&user_token,
|
||||
&mut transaction,
|
||||
).await.map_err(|err| {
|
||||
error!("Error creating artcraft wallet for user {:?}: {:?}", user_token, err);
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|err| {
|
||||
error!("Error committing MySQL transaction: {:?}", err);
|
||||
CommonWebError::ServerError
|
||||
})?;
|
||||
|
||||
Ok(Json(GetOrCreateSessionArtcraftWalletResponse {
|
||||
success: true,
|
||||
wallet_token,
|
||||
}))
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod get_or_create_session_artcraft_wallet_handler;
|
||||
+2
@@ -13,6 +13,7 @@ use crate::http_server::routes::application_routes::tag_routes::add_tag_routes;
|
||||
use crate::http_server::routes::application_routes::user_bookmarks_routes::add_user_bookmarks_routes;
|
||||
use crate::http_server::routes::application_routes::user_rating_routes::add_user_rating_routes;
|
||||
use crate::http_server::routes::application_routes::user_routes::add_user_routes;
|
||||
use crate::http_server::routes::application_routes::wallet_routes::add_wallet_routes;
|
||||
use crate::http_server::routes::application_routes::webhook_routes::add_webhook_routes;
|
||||
use crate::http_server::routes::application_routes::weights_routes::add_weights_routes;
|
||||
use actix_http::body::MessageBody;
|
||||
@@ -62,6 +63,7 @@ where
|
||||
app = add_user_routes(app); // /create_account, /session, /login, /logout, etc.
|
||||
|
||||
// Artcraft Billing pieces
|
||||
app = add_wallet_routes(app); // /v1/wallets/...
|
||||
app = add_credits_routes(app); // /v1/credits/...
|
||||
app = add_stripe_artcraft_routes(app); // /v1/stripe_artcraft/...
|
||||
app = add_subscription_routes(app); // /v1/subscriptions/...
|
||||
|
||||
+5
@@ -29,6 +29,7 @@ use crate::http_server::endpoints::generate::video::generate_veo_3_fast_image_to
|
||||
use crate::http_server::endpoints::generate::video::generate_veo_3_image_to_video_handler::generate_veo_3_image_to_video_handler;
|
||||
use crate::http_server::endpoints::generate::video::multi_function::kling_2p5_turbo_pro_multi_function_video_gen_handler::kling_2p5_turbo_pro_multi_function_video_gen_handler;
|
||||
use crate::http_server::endpoints::generate::video::multi_function::kling_2p6_pro_multi_function_video_gen_handler::kling_2p6_pro_multi_function_video_gen_handler;
|
||||
use crate::http_server::endpoints::generate::video::multi_function::seedance_2p0_multi_function_video_gen_handler::seedance_2p0_multi_function_video_gen_handler;
|
||||
use crate::http_server::endpoints::generate::video::multi_function::sora_2_multi_function_video_gen_handler::sora_2_multi_function_video_gen_handler;
|
||||
use crate::http_server::endpoints::generate::video::multi_function::sora_2_pro_multi_function_video_gen_handler::sora_2_pro_multi_function_video_gen_handler;
|
||||
use crate::http_server::endpoints::generate::video::multi_function::veo_3p1_fast_multi_function_video_gen_handler::veo_3p1_fast_multi_function_video_gen_handler;
|
||||
@@ -140,6 +141,10 @@ where
|
||||
.route(web::post().to(kling_2p6_pro_multi_function_video_gen_handler))
|
||||
.route(web::head().to(|| HttpResponse::Ok()))
|
||||
)
|
||||
.service(web::resource("/seedance_2p0")
|
||||
.route(web::post().to(seedance_2p0_multi_function_video_gen_handler))
|
||||
.route(web::head().to(|| HttpResponse::Ok()))
|
||||
)
|
||||
.service(web::resource("/sora_2")
|
||||
.route(web::post().to(sora_2_multi_function_video_gen_handler))
|
||||
.route(web::head().to(|| HttpResponse::Ok()))
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
mod analytics_routes;
|
||||
mod comments_routes;
|
||||
mod credits_routes;
|
||||
mod featured_item_routes;
|
||||
@@ -5,14 +6,14 @@ mod generate_routes;
|
||||
mod job_routes;
|
||||
mod media_files_routes;
|
||||
mod moderation_routes;
|
||||
mod prompts_routes;
|
||||
mod stripe_artcraft_routes;
|
||||
mod subscription_routes;
|
||||
mod tag_routes;
|
||||
mod user_bookmarks_routes;
|
||||
mod user_rating_routes;
|
||||
mod user_routes;
|
||||
mod wallet_routes;
|
||||
mod webhook_routes;
|
||||
mod weights_routes;
|
||||
pub (super) mod add_application_routes;
|
||||
mod prompts_routes;
|
||||
mod analytics_routes;
|
||||
pub (super) mod add_application_routes;
|
||||
+24
@@ -0,0 +1,24 @@
|
||||
use crate::http_server::endpoints::user_wallets::get_or_create_session_artcraft_wallet_handler::get_or_create_session_artcraft_wallet_handler;
|
||||
use actix_http::body::MessageBody;
|
||||
use actix_service::ServiceFactory;
|
||||
use actix_web::dev::{ServiceRequest, ServiceResponse};
|
||||
use actix_web::{web, App, Error, HttpResponse};
|
||||
|
||||
pub fn add_wallet_routes<T, B>(app: App<T>) -> App<T>
|
||||
where
|
||||
B: MessageBody,
|
||||
T: ServiceFactory<
|
||||
ServiceRequest,
|
||||
Config = (),
|
||||
Response = ServiceResponse<B>,
|
||||
Error = Error,
|
||||
InitError = (),
|
||||
>,
|
||||
{
|
||||
app.service(web::scope("/v1/wallets")
|
||||
.service(web::resource("/session_artcraft_wallet")
|
||||
.route(web::get().to(get_or_create_session_artcraft_wallet_handler))
|
||||
.route(web::head().to(|| HttpResponse::Ok()))
|
||||
)
|
||||
)
|
||||
}
|
||||
@@ -80,7 +80,7 @@ use crate::http_server::web_utils::handle_multipart_error::handle_multipart_erro
|
||||
use crate::http_server::web_utils::scoped_temp_dir_creator::ScopedTempDirCreator;
|
||||
use crate::state::certs::google_sign_in_cert::GoogleSignInCert;
|
||||
use crate::state::memory_cache::model_token_to_info_cache::ModelTokenToInfoCache;
|
||||
use crate::state::server_state::{DurableInMemoryCaches, EnvConfig, EphemeralInMemoryCaches, FalData, InMemoryCaches, OpenAiData, ResendData, ServerInfo, ServerState, StaticFeatureFlags, StripeSettings, TrollBans};
|
||||
use crate::state::server_state::{DurableInMemoryCaches, EnvConfig, EphemeralInMemoryCaches, FalData, InMemoryCaches, OpenAiData, ResendData, Seedance2ProData, ServerInfo, ServerState, StaticFeatureFlags, StripeSettings, TrollBans};
|
||||
use crate::threads::db_health_checker_thread::db_health_check_status::HealthCheckStatus;
|
||||
use crate::threads::db_health_checker_thread::db_health_checker_thread::db_health_checker_thread;
|
||||
use crate::threads::poll_ip_banlist_thread::poll_ip_bans;
|
||||
@@ -124,6 +124,13 @@ const LOG_FORMAT : &str =
|
||||
#[actix_web::main]
|
||||
async fn main() -> AnyhowResult<()> {
|
||||
|
||||
// NB: Both ring (via reqwest 0.12/hyper-rustls) and aws-lc-rs (via quinn/resend-rs)
|
||||
// are compiled into this binary as rustls crypto providers. rustls 0.23 panics at
|
||||
// runtime if it can't auto-select a single provider, so we install ring explicitly.
|
||||
rustls::crypto::ring::default_provider()
|
||||
.install_default()
|
||||
.expect("Failed to install rustls crypto provider");
|
||||
|
||||
let container_environment = bootstrap(BootstrapArgs {
|
||||
app_name: "storyteller-web",
|
||||
default_logging_override: Some(DEFAULT_RUST_LOG),
|
||||
@@ -400,6 +407,8 @@ async fn main() -> AnyhowResult<()> {
|
||||
let fal_api_key = FalApiKey::new(easyenv::get_env_string_required("FAL_API_KEY")?);
|
||||
let fal_webhook_url = easyenv::get_env_string_required("FAL_WEBHOOK_URL")?;
|
||||
|
||||
let seedance2pro_cookies = easyenv::get_env_string_required("SEEDANCE2PRO_COOKIES")?;
|
||||
|
||||
let openai_api_key= easyenv::get_env_string_required("OPENAI_API_KEY")?;
|
||||
|
||||
let resend_api_key = easyenv::get_env_string_required("RESEND_API_KEY")?;
|
||||
@@ -458,7 +467,10 @@ async fn main() -> AnyhowResult<()> {
|
||||
fal: FalData {
|
||||
api_key: fal_api_key,
|
||||
webhook_url: fal_webhook_url,
|
||||
},
|
||||
},
|
||||
seedance2pro: Seedance2ProData {
|
||||
cookies: seedance2pro_cookies,
|
||||
},
|
||||
openai: OpenAiData {
|
||||
api_key: openai_api_key,
|
||||
},
|
||||
|
||||
@@ -92,7 +92,9 @@ pub struct ServerState {
|
||||
pub auto_gc_bucket_client: BucketClient,
|
||||
|
||||
pub fal: FalData,
|
||||
|
||||
|
||||
pub seedance2pro: Seedance2ProData,
|
||||
|
||||
pub openai: OpenAiData,
|
||||
|
||||
pub resend: ResendData,
|
||||
@@ -311,6 +313,13 @@ pub struct FalData {
|
||||
pub webhook_url: String,
|
||||
}
|
||||
|
||||
/// Seedance 2 Pro integration
|
||||
#[derive(Clone)]
|
||||
pub struct Seedance2ProData {
|
||||
/// Cookie string for seedance2-pro.com
|
||||
pub cookies: String,
|
||||
}
|
||||
|
||||
/// OpenAI integration
|
||||
#[derive(Clone)]
|
||||
pub struct OpenAiData {
|
||||
|
||||
@@ -1,9 +1,14 @@
|
||||
import { useState, useRef, useEffect, useCallback } from "react";
|
||||
import { useState, useRef, useEffect, useCallback, useMemo } from "react";
|
||||
import { useShallow } from "zustand/react/shallow";
|
||||
import { DRAW_LAYER_ID, INPAINT_LAYER_ID, PaintSurface } from "./PaintSurface";
|
||||
import "./App.css";
|
||||
import PromptEditor from "./PromptEditor/PromptEditor";
|
||||
import SideToolbar from "./components/ui/SideToolbar";
|
||||
import { AspectRatioType, useSceneStore } from "./stores/SceneState";
|
||||
import {
|
||||
AspectRatioType,
|
||||
type SceneState,
|
||||
useSceneStore,
|
||||
} from "./stores/SceneState";
|
||||
import { useUndoRedoHotkeys } from "./hooks/useUndoRedoHotkeys";
|
||||
import { useDeleteHotkeys } from "./hooks/useDeleteHotkeys";
|
||||
import { useCopyPasteHotkeys } from "./hooks/useCopyPasteHotkeys";
|
||||
@@ -18,16 +23,26 @@ import {
|
||||
useSelectedImageModel,
|
||||
useSelectedProviderForModel,
|
||||
} from "@storyteller/ui-model-selector";
|
||||
import { EnqueueEditImage, EnqueueEditImageRequest, EnqueueEditImageResolution, EnqueueEditImageSize, EnqueueImageInpaint, EnqueueImageInpaintRequest, useCanvasBgRemovedEvent } from "@storyteller/tauri-api";
|
||||
import {
|
||||
EnqueueEditImage,
|
||||
EnqueueEditImageRequest,
|
||||
EnqueueEditImageResolution,
|
||||
EnqueueEditImageSize,
|
||||
EnqueueImageInpaint,
|
||||
EnqueueImageInpaintRequest,
|
||||
useCanvasBgRemovedEvent,
|
||||
} from "@storyteller/tauri-api";
|
||||
import { HelpMenuButton } from "@storyteller/ui-help-menu";
|
||||
import { GenerationProvider } from "@storyteller/api-enums";
|
||||
import { HistoryStack, ImageBundle } from "../PageEdit/HistoryStack";
|
||||
import { BaseImageSelector, BaseSelectorImage } from "../PageEdit/BaseImageSelector";
|
||||
import {
|
||||
BaseImageSelector,
|
||||
BaseSelectorImage,
|
||||
} from "../PageEdit/BaseImageSelector";
|
||||
import { normalizeCanvas } from "~/Helpers/CanvasHelpers";
|
||||
import { EncodeImageBitmapToBase64 } from "./utilities/EncodeImageBitmapToBase64";
|
||||
import { RefImage, usePrompt2DStore } from "@storyteller/ui-promptbox";
|
||||
import { PromptsApi } from "@storyteller/api";
|
||||
import toast from "react-hot-toast";
|
||||
|
||||
const PAGE_ID: ModelPage = ModelPage.Canvas2D;
|
||||
|
||||
@@ -62,26 +77,123 @@ const PageDraw = () => {
|
||||
const [isSelecting, setIsSelecting] = useState<boolean>(false);
|
||||
const stageRef = useRef<Konva.Stage>({} as Konva.Stage);
|
||||
const transformerRefs = useRef<{ [key: string]: Konva.Transformer }>({});
|
||||
const store = useSceneStore();
|
||||
|
||||
/*
|
||||
* Scene store: use a selector + useShallow to avoid re-rendering on every store change.
|
||||
* Without this, useSceneStore() subscribes to the whole store (e.g. moveNode during
|
||||
* drag updates state constantly) and PageDraw would re-render every frame.
|
||||
*
|
||||
* useShallow compares the selected object by top-level keys: we only re-render when
|
||||
* one of these values actually changes. Do NOT add:
|
||||
* - cursorPosition / cursorVisible: updated every mouse move → would re-render constantly.
|
||||
* - historyImageNodeMap: mutated in place in the store (reference never changes),
|
||||
* so shallow would never see updates; fix the store to replace the Map if you need it.
|
||||
* When adding new store fields here, prefer primitive/array/object refs that the store
|
||||
* replaces (e.g. new array from .map()), not in-place mutations.
|
||||
*
|
||||
* IMPORTANT: The selector function is memoized to prevent infinite loops. If you modify
|
||||
* this selector, ensure the function reference stays stable (use useMemo if needed).
|
||||
*/
|
||||
const selector = useMemo(
|
||||
() => (state: SceneState) => ({
|
||||
baseImageInfo: state.baseImageInfo,
|
||||
baseImageBitmap: state.baseImageBitmap,
|
||||
nodes: state.nodes,
|
||||
selectedNodeIds: state.selectedNodeIds,
|
||||
lineNodes: state.lineNodes,
|
||||
activeTool: state.activeTool,
|
||||
currentShape: state.currentShape,
|
||||
fillColor: state.fillColor,
|
||||
brushColor: state.brushColor,
|
||||
brushSize: state.brushSize,
|
||||
historyImageBundles: state.historyImageBundles,
|
||||
getAspectRatioDimensions: state.getAspectRatioDimensions,
|
||||
finishRemoveBackground: state.finishRemoveBackground,
|
||||
createImageFromUrl: state.createImageFromUrl,
|
||||
createImageFromFile: state.createImageFromFile,
|
||||
setBaseImageInfo: state.setBaseImageInfo,
|
||||
RESET: state.RESET,
|
||||
clearLineNodes: state.clearLineNodes,
|
||||
setNodes: state.setNodes,
|
||||
removeHistoryImage: state.removeHistoryImage,
|
||||
addHistoryImageBundle: state.addHistoryImageBundle,
|
||||
setAspectRatioType: state.setAspectRatioType,
|
||||
setActiveTool: state.setActiveTool,
|
||||
selectNode: state.selectNode,
|
||||
setCurrentShape: state.setCurrentShape,
|
||||
setBrushColor: state.setBrushColor,
|
||||
setBrushSize: state.setBrushSize,
|
||||
setBrushOpacity: state.setBrushOpacity,
|
||||
setFillColor: state.setFillColor,
|
||||
toggleLock: state.toggleLock,
|
||||
beginRemoveBackground: state.beginRemoveBackground,
|
||||
bringToFront: state.bringToFront,
|
||||
bringForward: state.bringForward,
|
||||
sendBackward: state.sendBackward,
|
||||
sendToBack: state.sendToBack,
|
||||
copySelectedItems: state.copySelectedItems,
|
||||
pasteItems: state.pasteItems,
|
||||
deleteSelectedItems: state.deleteSelectedItems,
|
||||
undo: state.undo,
|
||||
redo: state.redo,
|
||||
}),
|
||||
[],
|
||||
);
|
||||
|
||||
const {
|
||||
baseImageInfo,
|
||||
baseImageBitmap,
|
||||
nodes,
|
||||
selectedNodeIds,
|
||||
lineNodes,
|
||||
activeTool,
|
||||
currentShape,
|
||||
fillColor,
|
||||
brushColor,
|
||||
brushSize,
|
||||
historyImageBundles,
|
||||
getAspectRatioDimensions,
|
||||
finishRemoveBackground,
|
||||
createImageFromUrl,
|
||||
createImageFromFile,
|
||||
setBaseImageInfo,
|
||||
RESET,
|
||||
clearLineNodes,
|
||||
setNodes,
|
||||
removeHistoryImage,
|
||||
addHistoryImageBundle,
|
||||
setAspectRatioType,
|
||||
setActiveTool,
|
||||
selectNode,
|
||||
setCurrentShape,
|
||||
setBrushColor,
|
||||
setBrushSize,
|
||||
setBrushOpacity,
|
||||
setFillColor,
|
||||
toggleLock,
|
||||
beginRemoveBackground,
|
||||
bringToFront,
|
||||
bringForward,
|
||||
sendBackward,
|
||||
sendToBack,
|
||||
copySelectedItems,
|
||||
pasteItems,
|
||||
deleteSelectedItems,
|
||||
undo,
|
||||
redo,
|
||||
} = useSceneStore(useShallow(selector));
|
||||
|
||||
const promptStoreProvider = usePrompt2DStore;
|
||||
const generationCount = promptStoreProvider((state) => state.generationCount);
|
||||
const setGenerationCount = promptStoreProvider((state) => state.setGenerationCount);
|
||||
const useSystemPrompt = promptStoreProvider((state) => state.useSystemPrompt);
|
||||
const referenceImages = promptStoreProvider((state) => state.referenceImages);
|
||||
const prompt = promptStoreProvider((state) => state.prompt);
|
||||
const setGenerationCount = promptStoreProvider(
|
||||
(state) => state.setGenerationCount,
|
||||
);
|
||||
|
||||
const baseImageKonvaRef = useRef<Konva.Image>({} as Konva.Image);
|
||||
const baseImageUrl = store.baseImageInfo?.url;
|
||||
const baseImageUrl = baseImageInfo?.url;
|
||||
const [pendingGenerations, setPendingGenerations] = useState<
|
||||
{ id: string; count: number }[]
|
||||
>([]);
|
||||
const addHistoryImageBundle = useSceneStore(
|
||||
(state) => state.addHistoryImageBundle,
|
||||
);
|
||||
const removeHistoryImage = useSceneStore((state) => state.removeHistoryImage);
|
||||
const historyImageBundles = useSceneStore(
|
||||
(state) => state.historyImageBundles,
|
||||
);
|
||||
|
||||
const selectedImageModel: ImageModel | undefined =
|
||||
useSelectedImageModel(PAGE_ID);
|
||||
@@ -92,11 +204,11 @@ const PageDraw = () => {
|
||||
const supportsMaskedInpainting =
|
||||
selectedImageModel?.usesInpaintingMask ?? false;
|
||||
|
||||
useDeleteHotkeys({ onDelete: store.deleteSelectedItems });
|
||||
useUndoRedoHotkeys({ undo: store.undo, redo: store.redo });
|
||||
useDeleteHotkeys({ onDelete: deleteSelectedItems });
|
||||
useUndoRedoHotkeys({ undo, redo });
|
||||
useCopyPasteHotkeys({
|
||||
onCopy: store.copySelectedItems,
|
||||
onPaste: store.pasteItems,
|
||||
onCopy: copySelectedItems,
|
||||
onPaste: pasteItems,
|
||||
});
|
||||
|
||||
useCanvasBgRemovedEvent(async (event) => {
|
||||
@@ -113,24 +225,19 @@ const PageDraw = () => {
|
||||
// const file = new File([blob], "generated_image.png", {
|
||||
// type: blob.type,
|
||||
// });
|
||||
store.finishRemoveBackground(
|
||||
nodeId,
|
||||
event.media_token,
|
||||
event.image_cdn_url,
|
||||
);
|
||||
finishRemoveBackground(nodeId, event.media_token, event.image_cdn_url);
|
||||
});
|
||||
|
||||
// Create a function to use the left layer ref and download the bitmap from it
|
||||
const getMaskArrayBuffer = async (): Promise<Uint8Array> => {
|
||||
if (
|
||||
!stageRef.current ||
|
||||
!baseImageKonvaRef.current
|
||||
) {
|
||||
if (!stageRef.current || !baseImageKonvaRef.current) {
|
||||
console.error("Stage or left panel ref is not available");
|
||||
throw new Error("Stage or left panel or base image ref is not available");
|
||||
}
|
||||
|
||||
const layer = stageRef.current.getLayers().find((l) => l.id() === INPAINT_LAYER_ID)!;
|
||||
const layer = stageRef.current
|
||||
.getLayers()
|
||||
.find((l) => l.id() === INPAINT_LAYER_ID)!;
|
||||
|
||||
// Get the canvas area that's covered by the image/rectangle
|
||||
const rect = baseImageKonvaRef.current;
|
||||
@@ -195,7 +302,7 @@ const PageDraw = () => {
|
||||
|
||||
console.log("Creating image from URL:", imageUrl);
|
||||
|
||||
store.createImageFromUrl(stagePoint.x, stagePoint.y, imageUrl);
|
||||
createImageFromUrl(stagePoint.x, stagePoint.y, imageUrl);
|
||||
|
||||
console.log(
|
||||
`Created image "${item.label}" at stage position:`,
|
||||
@@ -214,12 +321,12 @@ const PageDraw = () => {
|
||||
handleGallery2DDrop as EventListener,
|
||||
);
|
||||
};
|
||||
}, [store]);
|
||||
// Stable action ref; no need to depend on full store.
|
||||
}, [createImageFromUrl]);
|
||||
|
||||
const handleImageUpload = async (files: File[]): Promise<void> => {
|
||||
// Determine current canvas dimensions from the store (according to aspect-ratio)
|
||||
const { width: canvasW, height: canvasH } =
|
||||
store.getAspectRatioDimensions();
|
||||
const { width: canvasW, height: canvasH } = getAspectRatioDimensions();
|
||||
|
||||
// Target maximum size – 85 % of the canvas in each direction
|
||||
const maxW = canvasW * 0.85;
|
||||
@@ -240,80 +347,20 @@ const PageDraw = () => {
|
||||
const x = (canvasW - finalW) / 2;
|
||||
const y = (canvasH - finalH) / 2;
|
||||
|
||||
store.createImageFromFile(x, y, file, finalW, finalH);
|
||||
createImageFromFile(x, y, file, finalW, finalH);
|
||||
};
|
||||
img.src = URL.createObjectURL(file);
|
||||
}
|
||||
};
|
||||
|
||||
const handleTauriEnqueue = async (image: ImageBitmap, aspectRatio: EnqueueEditImageSize | undefined, resolution: EnqueueEditImageResolution | undefined, subscriberId: string) => {
|
||||
if (image === undefined) {
|
||||
console.log("image is undefined");
|
||||
return;
|
||||
}
|
||||
|
||||
const api = new PromptsApi();
|
||||
const base64Bitmap = await EncodeImageBitmapToBase64(image);
|
||||
|
||||
const byteString = atob(base64Bitmap);
|
||||
const mimeString = "image/png";
|
||||
|
||||
const ab = new ArrayBuffer(byteString.length);
|
||||
const ia = new Uint8Array(ab);
|
||||
|
||||
for (let i = 0; i < byteString.length; i++) {
|
||||
ia[i] = byteString.charCodeAt(i);
|
||||
}
|
||||
|
||||
const uuid = crypto.randomUUID(); // Generate a new UUID
|
||||
const file = new File([ab], `${uuid}.png`, { type: mimeString });
|
||||
|
||||
const snapshotMediaToken = await api.uploadSceneSnapshot({
|
||||
screenshot: file,
|
||||
});
|
||||
|
||||
if (snapshotMediaToken.data === undefined) {
|
||||
toast.error("Error: Unable to upload scene snapshot Please try again.");
|
||||
return;
|
||||
}
|
||||
|
||||
console.log("useSystemPrompt", useSystemPrompt);
|
||||
console.log("Snapshot media token:", snapshotMediaToken.data);
|
||||
|
||||
const request: EnqueueEditImageRequest = {
|
||||
model: selectedImageModel,
|
||||
scene_image_media_token: snapshotMediaToken.data!,
|
||||
image_media_tokens: referenceImages
|
||||
.map((image) => image.mediaToken)
|
||||
.filter((t) => t.length > 0),
|
||||
disable_system_prompt: !useSystemPrompt,
|
||||
prompt: prompt,
|
||||
image_count: generationCount,
|
||||
aspect_ratio: aspectRatio,
|
||||
image_resolution: resolution,
|
||||
frontend_caller: "image_editor",
|
||||
frontend_subscriber_id: subscriberId,
|
||||
};
|
||||
|
||||
if (selectedProvider) {
|
||||
request.provider = selectedProvider;
|
||||
}
|
||||
|
||||
const generateResponse = await EnqueueEditImage(request);
|
||||
console.log("generateResponse", generateResponse);
|
||||
return generateResponse;
|
||||
}
|
||||
|
||||
const getCompositeCanvasFile = useCallback(async (): Promise<File | null> => {
|
||||
if (
|
||||
!stageRef.current ||
|
||||
!baseImageKonvaRef.current ||
|
||||
!store.baseImageBitmap
|
||||
) {
|
||||
if (!stageRef.current || !baseImageKonvaRef.current || !baseImageBitmap) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const editsLayer = stageRef.current.getLayers().find((l) => l.id() === DRAW_LAYER_ID);
|
||||
const editsLayer = stageRef.current
|
||||
.getLayers()
|
||||
.find((l) => l.id() === DRAW_LAYER_ID);
|
||||
|
||||
if (!editsLayer) {
|
||||
console.error("Edits layer not found");
|
||||
@@ -328,7 +375,7 @@ const PageDraw = () => {
|
||||
const ctx = canvas.getContext("2d");
|
||||
if (!ctx) return null;
|
||||
|
||||
ctx.drawImage(store.baseImageBitmap, 0, 0, width, height);
|
||||
ctx.drawImage(baseImageBitmap, 0, 0, width, height);
|
||||
|
||||
const markerLayerCanvas = editsLayer.toCanvas({
|
||||
x: stageRef.current.x(),
|
||||
@@ -347,7 +394,7 @@ const PageDraw = () => {
|
||||
const blob = await canvas.convertToBlob({ type: "image/png" });
|
||||
const uuid = crypto.randomUUID();
|
||||
return new File([blob], `${uuid}.png`, { type: "image/png" });
|
||||
}, [store.baseImageBitmap]);
|
||||
}, [baseImageBitmap]);
|
||||
|
||||
const handleGenerate = useCallback(
|
||||
async (
|
||||
@@ -359,7 +406,7 @@ const PageDraw = () => {
|
||||
selectedProvider?: GenerationProvider;
|
||||
},
|
||||
) => {
|
||||
const editedImageToken = store.baseImageInfo?.mediaToken;
|
||||
const editedImageToken = baseImageInfo?.mediaToken;
|
||||
|
||||
if (!editedImageToken) {
|
||||
console.error("Base image is not available");
|
||||
@@ -400,7 +447,7 @@ const PageDraw = () => {
|
||||
}
|
||||
};
|
||||
|
||||
const { width, height } = store.getAspectRatioDimensions();
|
||||
const { width, height } = getAspectRatioDimensions();
|
||||
const subscriberId: string =
|
||||
crypto?.randomUUID?.() ??
|
||||
`inpaint-${Date.now()}-${Math.random().toString(36).slice(2)}`;
|
||||
@@ -518,7 +565,9 @@ const PageDraw = () => {
|
||||
);
|
||||
throw error;
|
||||
}
|
||||
}, [generationCount, getCompositeCanvasFile, selectedImageModel, store]);
|
||||
},
|
||||
[generationCount, getCompositeCanvasFile, selectedImageModel],
|
||||
);
|
||||
|
||||
const onFitPressed = async () => {
|
||||
// Get the stage and its container dimensions
|
||||
@@ -530,8 +579,8 @@ const PageDraw = () => {
|
||||
const containerHeight = stage.container().offsetHeight;
|
||||
|
||||
// Get canvas dimensions from store aspect ratio
|
||||
const canvasW = store.getAspectRatioDimensions().width;
|
||||
const canvasH = store.getAspectRatioDimensions().height;
|
||||
const canvasW = getAspectRatioDimensions().width;
|
||||
const canvasH = getAspectRatioDimensions().height;
|
||||
|
||||
// Add padding to ensure canvas doesn't touch the edges
|
||||
const padding = 40;
|
||||
@@ -560,15 +609,17 @@ const PageDraw = () => {
|
||||
};
|
||||
|
||||
// When the model inpainting support changes, we need to auto-change the tool so it's not set to inpainting
|
||||
// Note: setActiveTool is a stable Zustand action, so we don't need it in deps
|
||||
useEffect(() => {
|
||||
if (!supportsMaskedInpainting && store.activeTool === "inpaint") {
|
||||
store.setActiveTool("select");
|
||||
if (!supportsMaskedInpainting && activeTool === "inpaint") {
|
||||
setActiveTool("select");
|
||||
}
|
||||
}, [store, supportsMaskedInpainting]);
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [activeTool, supportsMaskedInpainting]);
|
||||
|
||||
// Auto-fit canvas to screen on initial load
|
||||
useEffect(() => {
|
||||
if (!store.baseImageBitmap) {
|
||||
if (!baseImageBitmap) {
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -597,11 +648,11 @@ const PageDraw = () => {
|
||||
|
||||
autoFitCanvas();
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [store.baseImageBitmap]);
|
||||
}, [baseImageBitmap]);
|
||||
|
||||
// Display image selector on launch, otherwise hide it
|
||||
// Also show loading state if info is set but image is loading
|
||||
if (!store.baseImageInfo || !store.baseImageBitmap) {
|
||||
if (!baseImageInfo || !baseImageBitmap) {
|
||||
return (
|
||||
<div
|
||||
className={
|
||||
@@ -613,11 +664,9 @@ const PageDraw = () => {
|
||||
<BaseImageSelector
|
||||
onImageSelect={(image: BaseSelectorImage) => {
|
||||
addHistoryImageBundle({ images: [image] });
|
||||
store.setBaseImageInfo(image);
|
||||
setBaseImageInfo(image);
|
||||
}}
|
||||
showLoading={
|
||||
store.baseImageInfo !== null && store.baseImageInfo === null
|
||||
}
|
||||
showLoading={baseImageInfo !== null && baseImageBitmap === null}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
@@ -629,29 +678,28 @@ const PageDraw = () => {
|
||||
<>
|
||||
<div className="fixed inset-0 -z-10 bg-ui-background" />
|
||||
<div
|
||||
className={`preserve-aspect-ratio fixed right-4 top-1/2 z-10 -translate-y-1/2 transform ${isSelecting ? "pointer-events-none" : "pointer-events-auto"
|
||||
}`}
|
||||
className={`preserve-aspect-ratio fixed right-4 top-1/2 z-10 -translate-y-1/2 transform ${
|
||||
isSelecting ? "pointer-events-none" : "pointer-events-auto"
|
||||
}`}
|
||||
>
|
||||
<HistoryStack
|
||||
onClear={() => {
|
||||
store.RESET();
|
||||
RESET();
|
||||
setPendingGenerations([]);
|
||||
}}
|
||||
imageBundles={historyImageBundles}
|
||||
pendingPlaceholders={pendingGenerations}
|
||||
blurredBackgroundUrl={baseImageUrl}
|
||||
onImageSelect={(baseImage) => {
|
||||
store.clearLineNodes();
|
||||
store.setNodes([]);
|
||||
store.setBaseImageInfo(baseImage);
|
||||
setBaseImageInfo(baseImage);
|
||||
}}
|
||||
onImageRemove={(baseImage) => {
|
||||
if (
|
||||
pendingGenerations.length === 0 &&
|
||||
store.historyImageBundles.length === 1 &&
|
||||
store.historyImageBundles[0].images.length <= 1
|
||||
historyImageBundles.length === 1 &&
|
||||
historyImageBundles[0].images.length <= 1
|
||||
) {
|
||||
store.RESET();
|
||||
RESET();
|
||||
} else {
|
||||
removeHistoryImage(baseImage);
|
||||
}
|
||||
@@ -662,12 +710,13 @@ const PageDraw = () => {
|
||||
onResolvePending={(id: string) =>
|
||||
setPendingGenerations((prev) => prev.filter((p) => p.id !== id))
|
||||
}
|
||||
selectedImageToken={store.baseImageInfo?.mediaToken}
|
||||
selectedImageToken={baseImageInfo?.mediaToken}
|
||||
/>
|
||||
</div>
|
||||
<div
|
||||
className={`preserve-aspect-ratio fixed bottom-0 left-1/2 z-10 -translate-x-1/2 transform ${isSelecting ? "pointer-events-none" : "pointer-events-auto"
|
||||
}`}
|
||||
className={`preserve-aspect-ratio fixed bottom-0 left-1/2 z-10 -translate-x-1/2 transform ${
|
||||
isSelecting ? "pointer-events-none" : "pointer-events-auto"
|
||||
}`}
|
||||
>
|
||||
<PromptEditor
|
||||
onAspectRatioChange={async (ratio: string) => {
|
||||
@@ -685,7 +734,7 @@ const PageDraw = () => {
|
||||
};
|
||||
|
||||
const aspectRatioType = ratioToType(ratio);
|
||||
store.setAspectRatioType(aspectRatioType);
|
||||
setAspectRatioType(aspectRatioType);
|
||||
|
||||
await new Promise((resolve) => requestAnimationFrame(resolve));
|
||||
onFitPressed();
|
||||
@@ -705,28 +754,28 @@ const PageDraw = () => {
|
||||
<SideToolbar
|
||||
className="fixed left-0 top-1/2 z-10 -translate-y-1/2 transform"
|
||||
onSelect={(): void => {
|
||||
store.setActiveTool("select");
|
||||
setActiveTool("select");
|
||||
}}
|
||||
onActivateShapeTool={(
|
||||
shape: "rectangle" | "circle" | "triangle",
|
||||
): void => {
|
||||
store.selectNode(null);
|
||||
store.setCurrentShape(shape);
|
||||
store.setActiveTool("shape");
|
||||
store.selectNode(null);
|
||||
selectNode(null);
|
||||
setCurrentShape(shape);
|
||||
setActiveTool("shape");
|
||||
selectNode(null);
|
||||
}}
|
||||
onPaintBrush={(hex: string, size: number, opacity: number): void => {
|
||||
store.setActiveTool("draw");
|
||||
store.setBrushColor(hex);
|
||||
store.setBrushSize(size);
|
||||
store.setBrushOpacity(opacity);
|
||||
setActiveTool("draw");
|
||||
setBrushColor(hex);
|
||||
setBrushSize(size);
|
||||
setBrushOpacity(opacity);
|
||||
}}
|
||||
onCanvasBackground={(hex: string): void => {
|
||||
console.log("Canvas background activated", { color: hex });
|
||||
// Add background change logic here
|
||||
// TODO: minor bug needs to update the preview panel
|
||||
// Debounce also causes issues with real time color change.
|
||||
store.setFillColor(hex);
|
||||
setFillColor(hex);
|
||||
}}
|
||||
onUploadImage={(): void => {
|
||||
// Create input element dynamically like in PromptEditor
|
||||
@@ -765,14 +814,14 @@ const PageDraw = () => {
|
||||
input.click();
|
||||
}}
|
||||
supportsMaskTool={supportsMaskedInpainting}
|
||||
activeToolId={store.activeTool}
|
||||
currentShape={store.currentShape}
|
||||
activeToolId={activeTool}
|
||||
currentShape={currentShape}
|
||||
/>
|
||||
<div className="relative z-0">
|
||||
<ContextMenuContainer
|
||||
onAction={(e, action) => {
|
||||
if (action === "contextMenu") {
|
||||
const hasSelection = store.selectedNodeIds.length > 0;
|
||||
const hasSelection = selectedNodeIds.length > 0;
|
||||
if (hasSelection) {
|
||||
console.log("An item is selected.");
|
||||
return true;
|
||||
@@ -786,51 +835,54 @@ const PageDraw = () => {
|
||||
onMenuAction={async (action) => {
|
||||
switch (action) {
|
||||
case "LOCK":
|
||||
store.toggleLock(store.selectedNodeIds);
|
||||
toggleLock(selectedNodeIds);
|
||||
break;
|
||||
case "REMOVE_BACKGROUND":
|
||||
await store.beginRemoveBackground(store.selectedNodeIds);
|
||||
await beginRemoveBackground(selectedNodeIds);
|
||||
break;
|
||||
case "BRING_TO_FRONT":
|
||||
store.bringToFront(store.selectedNodeIds);
|
||||
bringToFront(selectedNodeIds);
|
||||
break;
|
||||
case "BRING_FORWARD":
|
||||
store.bringForward(store.selectedNodeIds);
|
||||
bringForward(selectedNodeIds);
|
||||
break;
|
||||
case "SEND_BACKWARD":
|
||||
store.sendBackward(store.selectedNodeIds);
|
||||
sendBackward(selectedNodeIds);
|
||||
break;
|
||||
case "SEND_TO_BACK":
|
||||
store.sendToBack(store.selectedNodeIds);
|
||||
sendToBack(selectedNodeIds);
|
||||
break;
|
||||
case "DUPLICATE":
|
||||
store.copySelectedItems();
|
||||
store.pasteItems();
|
||||
copySelectedItems();
|
||||
pasteItems();
|
||||
break;
|
||||
case "DELETE":
|
||||
store.deleteSelectedItems();
|
||||
deleteSelectedItems();
|
||||
break;
|
||||
default:
|
||||
// No action
|
||||
}
|
||||
}}
|
||||
isLocked={store.selectedNodeIds.some((id) => {
|
||||
const node = store.nodes.find((n) => n.id === id);
|
||||
const lineNode = store.lineNodes.find((n) => n.id === id);
|
||||
isLocked={selectedNodeIds.some((id: string) => {
|
||||
const node = nodes.find((n: (typeof nodes)[number]) => n.id === id);
|
||||
const lineNode = lineNodes.find(
|
||||
(n: (typeof lineNodes)[number]) => n.id === id,
|
||||
);
|
||||
return (node?.locked || lineNode?.locked) ?? false;
|
||||
})}
|
||||
>
|
||||
<PaintSurface
|
||||
nodes={store.nodes}
|
||||
selectedNodeIds={store.selectedNodeIds}
|
||||
nodes={nodes}
|
||||
lineNodes={lineNodes}
|
||||
selectedNodeIds={selectedNodeIds}
|
||||
onCanvasSizeChange={(width: number, height: number): void => {
|
||||
canvasWidth.current = width;
|
||||
canvasHeight.current = height;
|
||||
}}
|
||||
fillColor={store.fillColor}
|
||||
activeTool={store.activeTool}
|
||||
brushColor={store.brushColor}
|
||||
brushSize={store.brushSize}
|
||||
fillColor={fillColor}
|
||||
activeTool={activeTool}
|
||||
brushColor={brushColor}
|
||||
brushSize={brushSize}
|
||||
onSelectionChange={setIsSelecting}
|
||||
stageRef={stageRef}
|
||||
transformerRefs={transformerRefs}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,6 @@
|
||||
import { Vector2d } from 'konva/lib/types';
|
||||
import { useEffect } from 'react';
|
||||
import { DragState } from '~/pages/PageEdit/EditPaintSurface';
|
||||
import { Vector2d } from "konva/lib/types";
|
||||
import { useEffect } from "react";
|
||||
import { DragState } from "~/pages/PageEdit/EditPaintSurface";
|
||||
// This piece of code ensures that when you are selecting or dragging and go off the stage,
|
||||
// it releases the controls mouse drag.
|
||||
export const useGlobalMouseUp = (
|
||||
@@ -10,11 +10,12 @@ export const useGlobalMouseUp = (
|
||||
setIsSelecting: (value: boolean) => void,
|
||||
isSelectingRef: React.MutableRefObject<boolean>,
|
||||
setSelectionRect: (value: any) => void,
|
||||
setLastPoint: (value: { x: number, y: number } | null) => void,
|
||||
onSelectionChange?: (isSelecting: boolean) => void,
|
||||
onBeforeReset?: () => void,
|
||||
) => {
|
||||
useEffect(() => {
|
||||
const handleGlobalMouseUp = () => {
|
||||
onBeforeReset?.();
|
||||
// Always reset all states when mouse is released anywhere
|
||||
setIsDragging(undefined);
|
||||
setIsDrawing(false);
|
||||
@@ -23,16 +24,15 @@ export const useGlobalMouseUp = (
|
||||
isSelectingRef.current = false;
|
||||
onSelectionChange?.(false);
|
||||
setSelectionRect(null);
|
||||
setLastPoint(null);
|
||||
};
|
||||
|
||||
// Listen to mouseup on the entire window
|
||||
window.addEventListener('mouseup', handleGlobalMouseUp);
|
||||
window.addEventListener('touchend', handleGlobalMouseUp);
|
||||
window.addEventListener("mouseup", handleGlobalMouseUp);
|
||||
window.addEventListener("touchend", handleGlobalMouseUp);
|
||||
|
||||
return () => {
|
||||
window.removeEventListener('mouseup', handleGlobalMouseUp);
|
||||
window.removeEventListener('touchend', handleGlobalMouseUp);
|
||||
window.removeEventListener("mouseup", handleGlobalMouseUp);
|
||||
window.removeEventListener("touchend", handleGlobalMouseUp);
|
||||
};
|
||||
}, [onSelectionChange]);
|
||||
}
|
||||
}, [onSelectionChange, onBeforeReset]);
|
||||
};
|
||||
|
||||
@@ -86,11 +86,17 @@ type SerializedNodeData = {
|
||||
locked: boolean;
|
||||
};
|
||||
|
||||
interface SceneState {
|
||||
interface HistoryNodeData {
|
||||
nodes: Node[];
|
||||
lineNodes: LineNode[];
|
||||
}
|
||||
|
||||
export interface SceneState {
|
||||
// Nodes
|
||||
nodes: Node[];
|
||||
selectedNodeIds: string[];
|
||||
lineNodes: LineNode[]; // Add lineNodes to state
|
||||
historyImageNodeMap: Map<BaseSelectorImage, HistoryNodeData>;
|
||||
|
||||
// Clipboard
|
||||
clipboard: (Node | LineNode)[]; // To store copied items
|
||||
@@ -292,6 +298,7 @@ export const useSceneStore = create<SceneState>((set, get, store) => ({
|
||||
nodes: [],
|
||||
lineNodes: [],
|
||||
selectedNodeIds: [],
|
||||
historyImageNodeMap: new Map<BaseSelectorImage, HistoryNodeData>(),
|
||||
clipboard: [],
|
||||
history: [],
|
||||
historyIndex: -1,
|
||||
@@ -1587,15 +1594,40 @@ export const useSceneStore = create<SceneState>((set, get, store) => ({
|
||||
},
|
||||
|
||||
setBaseImageInfo: (image: BaseSelectorImage | null) => {
|
||||
set({ baseImageInfo: image });
|
||||
|
||||
// If the new base image is null, return early because nothing needs to be loaded.
|
||||
if (!image) {
|
||||
// Set the new base image info - setting null is allowed.
|
||||
set({ baseImageInfo: null });
|
||||
return;
|
||||
}
|
||||
|
||||
// Load the image bitmap as well
|
||||
// Save the current node data to the history image node map
|
||||
const currentBaseImage = get().baseImageInfo;
|
||||
if (currentBaseImage) {
|
||||
get().historyImageNodeMap.set(currentBaseImage, {
|
||||
nodes: get().nodes,
|
||||
lineNodes: get().lineNodes,
|
||||
});
|
||||
}
|
||||
|
||||
// Set the new base image info
|
||||
set({ baseImageInfo: image });
|
||||
|
||||
// Load the image bitmap now
|
||||
const imgBitmap = new Image();
|
||||
imgBitmap.onload = () => {
|
||||
// Once the base image is set, check previous node data - or create new empty one
|
||||
const previousNodeData = get().historyImageNodeMap.get(image);
|
||||
if (previousNodeData) {
|
||||
set({
|
||||
nodes: previousNodeData.nodes,
|
||||
lineNodes: previousNodeData.lineNodes,
|
||||
});
|
||||
} else {
|
||||
set({ nodes: [], lineNodes: [] });
|
||||
}
|
||||
|
||||
// Finally, set the base image bitmap
|
||||
set({ baseImageBitmap: imgBitmap });
|
||||
};
|
||||
imgBitmap.onerror = (event) => {
|
||||
@@ -1615,7 +1647,10 @@ export const useSceneStore = create<SceneState>((set, get, store) => ({
|
||||
},
|
||||
|
||||
clearHistoryImages: () => {
|
||||
set({ historyImageBundles: [] });
|
||||
set({
|
||||
historyImageBundles: [],
|
||||
historyImageNodeMap: new Map<BaseSelectorImage, HistoryNodeData>(),
|
||||
});
|
||||
},
|
||||
|
||||
addHistoryImageBundle(bundle) {
|
||||
@@ -1627,6 +1662,10 @@ export const useSceneStore = create<SceneState>((set, get, store) => ({
|
||||
removeHistoryImage(image) {
|
||||
console.log("Removing history image:", image);
|
||||
set((state) => {
|
||||
// Remove the history image node data
|
||||
state.historyImageNodeMap.delete(image);
|
||||
|
||||
// Remove the image from the history image bundles
|
||||
const updatedBundles = state.historyImageBundles
|
||||
.map((bundle) => {
|
||||
return {
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user