Compare commits

...
Author SHA1 Message Date
Tanish Manku 942005282a Some bug fixes and improvements for FK 2025-04-04 00:38:31 -07:00
Tanish Manku 14474e68c2 Add basic Forward Kinematics support (#475)
There's still some glitches around selection and deletion.
2025-04-04 00:34:09 -07:00
Michael Chung 2f44da8ac7 Changed api calls to use file upload and uuid idempotency token and stuff works. 2025-04-03 23:10:23 -04:00
Brandon Thomas 3d67f86a19 set sora creds externally 2025-04-03 22:26:38 -04:00
Brandon Thomas 2913d91d0d fix media file downloading 2025-04-03 18:47:57 -04:00
ArEnScandMichael Chung ceb77505e1 Feature/gen (#474)
Co-authored-by: Michael Chung <michaelchung@michaels-mbp-2.home>
2025-04-03 18:01:29 -04:00
ArEnScandMichael Chung 0cad3fe7da Wire screenshot upload and generate to the engine. (#470)
Co-authored-by: Michael Chung <michaelchung@michaels-mbp-2.home>
2025-04-03 17:38:34 -04:00
Tanish Manku 1ced74040f Fix: Letterboxing blocking pointer events 2025-04-03 14:09:35 -07:00
Brandon ThomasandKasisnu c9dd545ae1 Use OpenAI Sora GPT 4o image gen (#472)
Co-authored-by: Kasisnu <5044120+kasisnu@users.noreply.github.com>
2025-04-03 17:07:20 -04:00
bflatastic 8ff91690c1 bring back objects, skybox, creatures tabs 2025-04-03 23:24:14 +07:00
BFlat 97bbe0fe07 add image ui (#471) 2025-04-03 23:01:02 +07:00
BFlat 9361634d83 final fix ratio default to 3:2 (#469) 2025-04-03 12:41:40 +07:00
bflatastic 7cf0363555 fixed ratio in render 3:2 2025-04-03 12:32:45 +07:00
bflatastic c2ee3b0ce1 render default to 3:2 ratio 2025-04-03 12:30:57 +07:00
BFlat ca6fc13171 aspect ratio update (#468)
- added 3:2, 2:3 ratios (defaulted to 3:2 [1200 x 800 px])
- added letterbox accordingly
- aspect ratio buttons to show icons for each ratio
- popover to close when item selected
2025-04-03 12:23:54 +07:00
Brandon Thomas cb1ea1f828 fix CDN in local development 2025-04-02 18:48:44 -04:00
bflatastic 010e3be5dc make letterbox always visible 2025-04-02 22:59:59 +07:00
bflatastic d8dab1a66e fix trimming 2 2025-04-02 22:42:28 +07:00
bflatastic 97b236d04d paste trim fix 2025-04-02 22:40:56 +07:00
bflatastic 15fe59cc19 added enqueueing state to generate button, limited textarea height and updates 2025-04-02 22:36:53 +07:00
BFlat 89bec31540 removed video controls and add prompt box (#467) 2025-04-02 22:01:19 +07:00
ArEnScandMichael Chung 7568a0ef86 Fixes for the new 3d editor (#465)
Removed the login page and timeline and confusing UI components.
Extended scene to the bottom of the screen.

Co-authored-by: Michael Chung <michaelchung@michaels-mbp-2.home>
2025-04-01 19:11:42 -04:00
Brandon Thomas 6180222e15 netlify.toml for editor2d 2025-04-01 00:48:53 -04:00
Brandon Thomas 6d45ce29e2 Adding copy of realtime as editor2d for web deployment 2025-04-01 00:44:28 -04:00
Brandon Thomas 1da56d6d22 netlify build outputs dir 2025-04-01 00:06:22 -04:00
Brandon Thomas 91a1c00442 reserved names 2025-04-01 00:00:52 -04:00
Brandon Thomas 3db942f1f6 artcraft cors 2025-03-31 23:58:48 -04:00
Brandon Thomas 2e56291649 unblock rust docker builds 2025-03-31 21:41:54 -04:00
Brandon Thomas b31badc6d1 Editor3D - this directory is a fork of "studio gen 2" by Tanish.
We'll be using this for our image generation roadmap:

https://docs.google.com/document/d/1O9YI_iLwXkl2wAJJ23dFp3eHR-JzhlRHCBPD8psG-1E/edit

The original Studio Gen 2 is in the `frontend/studio-gen-2` directory.

The original Studio Gen 1 is in the `frontend/studio-gen-1` directory
2025-03-31 13:28:59 -04:00
Brandon Thomas ed399fcb21 rename storyteller animator to "studio-gen-2" 2025-03-31 13:25:51 -04:00
Brandon Thomas d858c9394b work in progress background removal adjustment 2025-03-28 01:33:52 -04:00
Brandon Thomas a19a6599ec rename model 2025-03-27 23:02:51 -04:00
Brandon Thomas 57c15f7c9f rename realtime to sketch 2025-03-27 22:41:40 -04:00
Brandon Thomas aa25ae8c87 load flux models from R2 2025-03-27 21:04:39 -04:00
Brandon Thomas c27055761e fix 2025-03-26 20:35:50 -04:00
Brandon Thomas abf3024892 add to downloaders 2025-03-26 20:35:11 -04:00
Brandon Thomas 2182d4dde1 lots of new models in R2 2025-03-26 20:32:05 -04:00
Brandon Thomas b632b927a3 save images to file 2025-03-26 13:12:54 -04:00
Brandon Thomas 4e2be12064 flux cherry picked in 2025-03-26 02:37:04 -04:00
Brandon Thomas 29b33cd327 save image handler 2025-03-25 22:43:25 -04:00
Brandon Thomas 032d0e1eae remove lock button 2025-03-25 20:05:49 -04:00
Brandon Thomas a3ddffeeee get rid of "transform" / "move" context menu 2025-03-25 19:54:51 -04:00
Brandon Thomas 28a5ff2f9f remove background from preview images 2025-03-25 19:43:48 -04:00
Brandon Thomas 1028e04e02 turn lcm println!() to debug!() level 2025-03-25 17:45:29 -04:00
Brandon Thomas 49c010db59 temporarily remove references to flux (temp fix) 2025-03-25 17:34:33 -04:00
Kasisnu 48ff60d3b8 super basic flux (#463) 2025-03-25 17:05:50 -04:00
Brandon Thomas 9c5b993c4f read seed from config too 2025-03-25 12:34:33 -04:00
Brandon Thomas a4aecc4d38 send seed as a parameter 2025-03-25 12:21:47 -04:00
Brandon Thomas 4d73230e84 add logging 2025-03-25 10:44:11 -04:00
Brandon Thomas a5a09f3e64 fixup inpaint stub 2025-03-25 01:44:20 -04:00
49894bd9b3 Generate Mode (#461)
Co-authored-by: Tanish Manku <-h>
Co-authored-by: bflatastic <bombayapirat@gmail.com>
2025-03-25 01:38:56 -04:00
3181 changed files with 155617 additions and 1529 deletions
Generated
+450 -110
View File
File diff suppressed because it is too large Load Diff
+22
View File
@@ -32,6 +32,7 @@ members = [
"crates/lib/api_clients/cloud_storage",
"crates/lib/api_clients/google_drive_common",
"crates/lib/api_clients/google_sign_in",
"crates/lib/api_clients/openai_sora_client",
"crates/lib/api_clients/thumbnail_generator",
"crates/lib/bootstrap",
"crates/lib/bucket_paths",
@@ -58,6 +59,7 @@ members = [
"crates/lib/files/mimetypes",
"crates/lib/files/videos",
"crates/lib/idempotency",
"crates/lib/image_algorithms",
"crates/lib/jobs_common",
"crates/lib/logging",
"crates/lib/markdown",
@@ -117,3 +119,23 @@ members = [
# see usage: https://crates.io/crates/cargo-hakari
"crates/my-workspace-hack",
]
[workspace.package]
publish = false
[workspace.dependencies]
# Hakari workspace hack (faster compile times)
my-workspace-hack = { path = "crates/my-workspace-hack" }
# Internal libs
errors = { path = "crates/lib/errors" }
# External libs
# Centrally pinned versions so that we can uniformly upgrade/downgrade versions across the workspace
# If a crate needs to use an older/newer version, the dependency should be managed in that crate's own Cargo.toml
anyhow = "1.0.97"
log = "0.4.27"
reqwest = { version = "0.12.15", features = ["json", "cookies"] }
serde = "1.0.219"
serde_derive = "1.0.219"
serde_json = "1.0.140"
@@ -40,6 +40,10 @@ pub enum InferenceJobType {
#[serde(rename = "studio_gen2")]
StudioGen2,
/// Sora GPT 4o image gen
#[serde(rename = "image_gen_api")]
ImageGenApi,
/// A job that turns "FBX" game engine files into "GLTF" files (Bevy-compatible).
#[serde(rename = "convert_fbx_gltf")]
ConvertFbxToGltf,
@@ -108,6 +112,7 @@ impl InferenceJobType {
Self::GptSovits => "gpt_sovits",
Self::ComfyUi => "comfy_ui",
Self::StudioGen2 => "studio_gen2",
Self::ImageGenApi => "image_gen_api",
Self::ConvertFbxToGltf => "convert_fbx_gltf",
Self::MocapNet => "mocap_net",
Self::RvcV2 => "rvc_v2",
@@ -132,6 +137,7 @@ impl InferenceJobType {
"gpt_sovits" => Ok(Self::GptSovits),
"comfy_ui" => Ok(Self::ComfyUi),
"studio_gen2" => Ok(Self::StudioGen2),
"image_gen_api" => Ok(Self::ImageGenApi),
"convert_fbx_gltf" => Ok(Self::ConvertFbxToGltf),
"mocap_net" => Ok(Self::MocapNet),
"rvc_v2" => Ok(Self::RvcV2),
@@ -159,6 +165,7 @@ impl InferenceJobType {
Self::GptSovits,
Self::ComfyUi,
Self::StudioGen2,
Self::ImageGenApi,
Self::ConvertFbxToGltf,
Self::MocapNet,
Self::RvcV2,
@@ -197,6 +204,7 @@ mod tests {
assert_serialization(InferenceJobType::GptSovits, "gpt_sovits");
assert_serialization(InferenceJobType::ComfyUi, "comfy_ui");
assert_serialization(InferenceJobType::StudioGen2, "studio_gen2");
assert_serialization(InferenceJobType::ImageGenApi, "image_gen_api");
assert_serialization(InferenceJobType::ConvertFbxToGltf, "convert_fbx_gltf");
assert_serialization(InferenceJobType::MocapNet, "mocap_net");
assert_serialization(InferenceJobType::RvcV2, "rvc_v2");
@@ -220,6 +228,7 @@ mod tests {
assert_eq!(InferenceJobType::GptSovits.to_str(), "gpt_sovits");
assert_eq!(InferenceJobType::ComfyUi.to_str(), "comfy_ui");
assert_eq!(InferenceJobType::StudioGen2.to_str(), "studio_gen2");
assert_eq!(InferenceJobType::ImageGenApi.to_str(), "image_gen_api");
assert_eq!(InferenceJobType::ConvertFbxToGltf.to_str(), "convert_fbx_gltf");
assert_eq!(InferenceJobType::MocapNet.to_str(), "mocap_net");
assert_eq!(InferenceJobType::RvcV2.to_str(), "rvc_v2");
@@ -243,6 +252,7 @@ mod tests {
assert_eq!(InferenceJobType::from_str("gpt_sovits").unwrap(), InferenceJobType::GptSovits);
assert_eq!(InferenceJobType::from_str("comfy_ui").unwrap(), InferenceJobType::ComfyUi);
assert_eq!(InferenceJobType::from_str("studio_gen2").unwrap(), InferenceJobType::StudioGen2);
assert_eq!(InferenceJobType::from_str("image_gen_api").unwrap(), InferenceJobType::ImageGenApi);
assert_eq!(InferenceJobType::from_str("convert_fbx_gltf").unwrap(), InferenceJobType::ConvertFbxToGltf);
assert_eq!(InferenceJobType::from_str("mocap_net").unwrap(), InferenceJobType::MocapNet);
assert_eq!(InferenceJobType::from_str("rvc_v2").unwrap(), InferenceJobType::RvcV2);
@@ -261,7 +271,7 @@ mod tests {
fn all_variants() {
// Static check
let mut variants = InferenceJobType::all_variants();
assert_eq!(variants.len(), 16);
assert_eq!(variants.len(), 17);
assert_eq!(variants.pop_first(), Some(InferenceJobType::VideoRender));
assert_eq!(variants.pop_first(), Some(InferenceJobType::LivePortrait));
assert_eq!(variants.pop_first(), Some(InferenceJobType::FaceFusion));
@@ -269,6 +279,7 @@ mod tests {
assert_eq!(variants.pop_first(), Some(InferenceJobType::GptSovits));
assert_eq!(variants.pop_first(), Some(InferenceJobType::ComfyUi));
assert_eq!(variants.pop_first(), Some(InferenceJobType::StudioGen2));
assert_eq!(variants.pop_first(), Some(InferenceJobType::ImageGenApi));
assert_eq!(variants.pop_first(), Some(InferenceJobType::ConvertFbxToGltf));
assert_eq!(variants.pop_first(), Some(InferenceJobType::MocapNet));
assert_eq!(variants.pop_first(), Some(InferenceJobType::RvcV2));
@@ -33,6 +33,10 @@ pub enum InferenceModelType {
#[serde(rename = "seed_vc")]
SeedVc,
/// NB: This is for Sora GPT 4o image gen
#[serde(rename = "image_gen_api")]
ImageGenApi,
#[serde(rename = "tacotron2")]
Tacotron2,
#[serde(rename = "vits")]
@@ -70,6 +74,7 @@ impl InferenceModelType {
Self::VallEX => "vall_e_x",
Self::RerenderAVideo => "rerender_a_video",
Self::StableDiffusion => "stable_diffusion",
Self::ImageGenApi => "image_gen_api",
Self::SeedVc => "seed_vc",
Self::MocapNet => "mocap_net",
Self::StyleTTS2 => "styletts2",
@@ -90,6 +95,7 @@ impl InferenceModelType {
"vall_e_x" => Ok(Self::VallEX),
"rerender_a_video" => Ok(Self::RerenderAVideo),
"stable_diffusion" => Ok(Self::StableDiffusion),
"image_gen_api" => Ok(Self::ImageGenApi),
"mocap_net" => Ok(Self::MocapNet),
"styletts2" => Ok(Self::StyleTTS2),
"comfy_ui" => Ok(Self::ComfyUi),
@@ -113,6 +119,7 @@ impl InferenceModelType {
Self::VallEX,
Self::RerenderAVideo,
Self::StableDiffusion,
Self::ImageGenApi,
Self::MocapNet,
Self::StyleTTS2,
Self::ConvertFbxToGltf,
@@ -140,6 +147,7 @@ mod tests {
assert_serialization(InferenceModelType::VallEX, "vall_e_x");
assert_serialization(InferenceModelType::RerenderAVideo, "rerender_a_video");
assert_serialization(InferenceModelType::StableDiffusion, "stable_diffusion");
assert_serialization(InferenceModelType::ImageGenApi, "image_gen_api");
assert_serialization(InferenceModelType::MocapNet, "mocap_net");
assert_serialization(InferenceModelType::ComfyUi, "comfy_ui");
assert_serialization(InferenceModelType::StyleTTS2, "styletts2");
@@ -159,6 +167,7 @@ mod tests {
assert_eq!(InferenceModelType::VallEX.to_str(), "vall_e_x");
assert_eq!(InferenceModelType::RerenderAVideo.to_str(), "rerender_a_video");
assert_eq!(InferenceModelType::StableDiffusion.to_str(), "stable_diffusion");
assert_eq!(InferenceModelType::ImageGenApi.to_str(), "image_gen_api");
assert_eq!(InferenceModelType::MocapNet.to_str(), "mocap_net");
assert_eq!(InferenceModelType::StyleTTS2.to_str(), "styletts2");
assert_eq!(InferenceModelType::ConvertFbxToGltf.to_str(), "convert_fbx_gltf");
@@ -176,6 +185,7 @@ mod tests {
assert_eq!(InferenceModelType::from_str("vall_e_x").unwrap(), InferenceModelType::VallEX);
assert_eq!(InferenceModelType::from_str("rerender_a_video").unwrap(), InferenceModelType::RerenderAVideo);
assert_eq!(InferenceModelType::from_str("stable_diffusion").unwrap(), InferenceModelType::StableDiffusion);
assert_eq!(InferenceModelType::from_str("image_gen_api").unwrap(), InferenceModelType::ImageGenApi);
assert_eq!(InferenceModelType::from_str("mocap_net").unwrap(), InferenceModelType::MocapNet);
assert_eq!(InferenceModelType::from_str("styletts2").unwrap(), InferenceModelType::StyleTTS2);
assert_eq!(InferenceModelType::from_str("comfy_ui").unwrap(), InferenceModelType::ComfyUi);
@@ -187,17 +197,18 @@ mod tests {
fn all_variants() {
// Static check
let mut variants = InferenceModelType::all_variants();
assert_eq!(variants.len(), 13);
assert_eq!(variants.len(), 15);
assert_eq!(variants.pop_first(), Some(InferenceModelType::ComfyUi));
assert_eq!(variants.pop_first(), Some(InferenceModelType::RvcV2));
assert_eq!(variants.pop_first(), Some(InferenceModelType::SadTalker));
assert_eq!(variants.pop_first(), Some(InferenceModelType::SeedVc));
assert_eq!(variants.pop_first(), Some(InferenceModelType::SoVitsSvc));
assert_eq!(variants.pop_first(), Some(InferenceModelType::SeedVc));
assert_eq!(variants.pop_first(), Some(InferenceModelType::Tacotron2));
assert_eq!(variants.pop_first(), Some(InferenceModelType::Vits));
assert_eq!(variants.pop_first(), Some(InferenceModelType::VallEX));
assert_eq!(variants.pop_first(), Some(InferenceModelType::RerenderAVideo));
assert_eq!(variants.pop_first(), Some(InferenceModelType::StableDiffusion));
assert_eq!(variants.pop_first(), Some(InferenceModelType::ImageGenApi));
assert_eq!(variants.pop_first(), Some(InferenceModelType::MocapNet));
assert_eq!(variants.pop_first(), Some(InferenceModelType::StyleTTS2));
assert_eq!(variants.pop_first(), Some(InferenceModelType::ConvertFbxToGltf));
@@ -91,6 +91,10 @@ pub enum PublicInferenceJobType {
/// Re-render a video is a video style transfer algorithm. We developed code
/// around it, but chose to develop AnimateDiff / ComfyUI support instead.
RerenderAVideo,
/// Sora GPT 4o image gen
#[serde(rename = "image_gen_api")]
ImageGenApi,
}
/// NB: Legacy API for older code.
@@ -119,6 +123,7 @@ impl PublicInferenceJobType {
InferenceJobType::Unknown => Self::Unknown,
InferenceJobType::BevyToWorkflow => Self::BevyToWorkflow,
InferenceJobType::RerenderAVideo => Self::RerenderAVideo,
InferenceJobType::ImageGenApi => Self::ImageGenApi,
}
}
@@ -140,6 +145,7 @@ impl PublicInferenceJobType {
Self::SeedVc => InferenceJobType::SeedVc,
Self::SoVitsSvc => InferenceJobType::SoVitsSvc,
Self::StableDiffusion => InferenceJobType::StableDiffusion,
Self::ImageGenApi => InferenceJobType::ImageGenApi,
Self::StyleTTS2 => InferenceJobType::StyleTTS2,
Self::Tacotron2 => InferenceJobType::Tacotron2,
Self::Unknown => InferenceJobType::Unknown,
+1
View File
@@ -20,6 +20,7 @@ normal = ["my-workspace-hack"]
# Internal
filesys = { path = "../../lib/files/filesys" }
image_algorithms = { path = "../../lib/image_algorithms" }
memory_caching = { path = "../../lib/caching/memory_caching" }
ml_weights_registry = { path = "../ml_weights_registry" }
primitives = { path = "../../lib/data/primitives" }
@@ -3,6 +3,7 @@ use image::{imageops, DynamicImage};
use ndarray::{Array, CowArray};
use ort::Value;
use std::path::Path;
use image_algorithms::remove_image_border::remove_image_border;
pub async fn remove_image_background<P: AsRef<Path>>(model_path: P, image: DynamicImage) -> anyhow::Result<DynamicImage> {
// NB: Real time needed cuDNN 9, but this uses cuDNN 8:
@@ -11,6 +12,10 @@ pub async fn remove_image_background<P: AsRef<Path>>(model_path: P, image: Dynam
// update-alternatives: using /usr/include/x86_64-linux-gnu/cudnn_v8.h to provide /usr/include/cudnn.h (libcudnn) in manual mode
let session = onnx_session(model_path.as_ref())?;
let image = process_dynamic_image(&session, image)?;
// Crop transparent border cruft away.
let image = remove_image_border(&image)?;
Ok(image)
}
@@ -0,0 +1,297 @@
use std::path::Path;
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
use crate::ml::model_cache::ModelCache;
use crate::ml::prompt_cache::PromptCache;
use candle_nn::VarBuilder;
use crate::ml::model_config::ModelConfig;
use anyhow::{anyhow, Error as E, Result};
use candle_core::{DType, IndexOp, Tensor, D, Device, Module};
use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
use candle_transformers::models::stable_diffusion::lcm::LCMScheduler;
use image::DynamicImage;
use log::{debug, info};
use rand::Rng;
use tauri::{AppHandle, Emitter};
use candle_transformers::models::{clip, t5, flux};
use tokenizers::Tokenizer;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FluxModel {
Schnell,
Dev,
}
pub struct FluxPipelineArgs<'a> {
pub image: &'a DynamicImage,
pub prompt: String,
pub model_cache: &'a ModelCache,
pub prompt_cache: &'a PromptCache,
pub app: &'a AppHandle,
pub flux_model: FluxModel,
pub active_blocks: Option<usize>,
pub use_cuda_stream: Option<bool>,
pub prefetch_next_batch: Option<bool>,
pub use_full_gpu: Option<bool>,
pub img2img_strength: Option<f64>,
pub seed: Option<u64>,
pub use_quantized_model: Option<bool>,
pub model_config: &'a ModelConfig,
pub flux_dev_or_schnell_path: &'a Path,
pub flux_schnell_quantized_gguf: &'a Path,
pub flux_autoencoder_path: &'a Path,
pub google_t5_model_path: &'a Path,
pub google_t5_config_path: &'a Path,
pub lmz_t5_tokenizer_path: &'a Path,
pub openai_clip_model_path: &'a Path,
pub openai_clip_tokenizer_json_path: &'a Path,
}
pub fn flux_pipeline(args: FluxPipelineArgs<'_>) -> Result<RgbImage> {
let FluxPipelineArgs {
prompt,
model_cache,
prompt_cache,
app,
image,
flux_model,
active_blocks,
use_cuda_stream,
prefetch_next_batch,
use_full_gpu,
img2img_strength,
seed,
use_quantized_model,
model_config ,
flux_dev_or_schnell_path,
flux_schnell_quantized_gguf,
flux_autoencoder_path,
google_t5_model_path,
google_t5_config_path,
lmz_t5_tokenizer_path,
openai_clip_model_path,
openai_clip_tokenizer_json_path,
} = args;
if use_quantized_model.unwrap_or(false) {
#[cfg(feature = "cuda")]
candle_core::quantized::cuda::set_force_dmmv(false);
}
let height = 1024;
let width = 1024;
let quantized = use_quantized_model.unwrap_or(false);
let device = model_config.device.clone();
let dtype = device.bf16_default_to_f32();
info!("Starting image generation with FLUX model:");
info!(" Model: {:?}", flux_model);
info!(" Prompt: {}", prompt);
info!(" Dimensions: {}x{}", width, height);
info!(" Device: {:?}", model_config.device);
let seed_value = seed.unwrap_or_else(|| rand::thread_rng().gen());
device.set_seed(seed_value)?;
info!("Using seed: {}", seed_value);
// Setup FLUX model configuration
let flux_config = match flux_model {
FluxModel::Dev => flux::model::Config::dev(),
FluxModel::Schnell => flux::model::Config::schnell(),
};
// Get T5 embeddings
let t5_emb = {
info!("Generating T5 embeddings for prompt...");
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[google_t5_model_path], dtype , &device)? };
let config = std::fs::read_to_string(google_t5_config_path)?;
let config: t5::Config = serde_json::from_str(&config)?;
let mut model = t5::T5EncoderModel::load(vb, &config)?;
let tokenizer = Tokenizer::from_file(lmz_t5_tokenizer_path).map_err(E::msg)?;
let mut tokens = tokenizer.encode(prompt.as_str(), true).map_err(E::msg)?.get_ids().to_vec();
tokens.resize(256, 0);
let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
model.forward(&input_token_ids)?
};
info!("T5 embeddings generated");
// Get CLIP embeddings
let clip_emb = {
info!("Generating CLIP embeddings for prompt...");
let vb =
unsafe { VarBuilder::from_mmaped_safetensors(&[openai_clip_model_path], dtype, &device)? };
// https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json
let config = clip::text_model::ClipTextConfig {
vocab_size: 49408,
projection_dim: 768,
activation: clip::text_model::Activation::QuickGelu,
intermediate_size: 3072,
embed_dim: 768,
max_position_embeddings: 77,
pad_with: None,
num_hidden_layers: 12,
num_attention_heads: 12,
};
let model =
clip::text_model::ClipTextTransformer::new(vb.pp("text_model"), &config)?;
let tokenizer = Tokenizer::from_file(openai_clip_tokenizer_json_path).map_err(E::msg)?;
let tokens = tokenizer
.encode(prompt.as_str(), true)
.map_err(E::msg)?
.get_ids()
.to_vec();
let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
model.forward(&input_token_ids)?
};
info!("CLIP embeddings generated");
// Process input image to tensor
info!("Converting input image to tensor...");
let img_tensor = dynamic_image_to_tensor(image, &device, dtype)?;
info!("Image tensor shape: {:?}", img_tensor.shape());
// Create initial noise - this is important to match the expected format
info!("Generating initial noise for FLUX...");
let initial_tensor = flux::sampling::get_noise(1, height, width, &device)?.to_dtype(dtype)?;
info!("Initial tensor shape: {:?}", initial_tensor.shape());
// Blend initial noise with input image if img2img_strength is specified
let initial_tensor = if let Some(strength) = img2img_strength {
if strength > 0.0 && strength < 1.0 {
info!("Using img2img with strength: {}", strength);
// Reshape the input tensor to match the noise tensor dimensions
let adjusted_img = if img_tensor.shape() != initial_tensor.shape() {
info!("Reshaping input image to match noise dimensions...");
// You may need more complex reshaping logic here depending on your use case
let (_, c, h, w) = initial_tensor.dims4()?;
img_tensor.upsample_nearest2d(h, w)?
} else {
img_tensor
};
// Blend the noise and the image according to strength
let strength_tensor = Tensor::new(strength as f32, &device)?;
let inverse_strength = Tensor::new(1.0 - strength as f32, &device)?;
let weighted_noise = (&initial_tensor * &strength_tensor)?;
let weighted_img = (&adjusted_img * &inverse_strength)?;
(&weighted_noise + &weighted_img)?
} else {
initial_tensor
}
} else {
initial_tensor
};
info!("Final initial tensor shape: {:?}", initial_tensor.shape());
// Create FLUX state
info!("Creating FLUX state...");
let state = if quantized {
flux::sampling::State::new(
&t5_emb.to_dtype(DType::F32)?,
&clip_emb.to_dtype(DType::F32)?,
&initial_tensor.to_dtype(DType::F32)?,
)?
} else {
flux::sampling::State::new(&t5_emb, &clip_emb, &initial_tensor)?
};
info!("{state:?}");
// Set appropriate timesteps based on model
let timesteps = match flux_model {
FluxModel::Dev => flux::sampling::get_schedule(28, Some((state.img.dim(1)?, 0.5, 1.15))),
FluxModel::Schnell => flux::sampling::get_schedule(4, None),
};
info!("Starting FLUX denoising with timesteps: {:?}", timesteps);
// Actual model loading and inference
let img_latent = if quantized {
// Quantized model path
info!("Loading quantized FLUX model...");
let model_file = match flux_model {
FluxModel::Schnell => flux_schnell_quantized_gguf,
FluxModel::Dev => return Err(anyhow!("Quantized model not available for FLUX.1-dev")),
};
info!("Loading quantized model from {}", model_file.to_str().unwrap());
let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(model_file, &device)?;
let mut model = flux::quantized_model::Flux::new(&flux_config, vb)?;
info!("Quantized model loaded on {:?}", device);
// Run denoising
info!("Running denoising...");
let img_latent = flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, &timesteps, 4.)?.to_dtype(dtype)?;
info!("Denoising completed");
img_latent
} else {
// Full precision model
info!("Loading FLUX model...");
let mut model = if use_full_gpu.unwrap_or(false) {
// Load entire model on GPU
info!("Loading entire model directly to GPU...");
let gpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[flux_dev_or_schnell_path], dtype, &device)? };
let mut model = flux::model::Flux::new(&flux_config, gpu_vb)?;
model.set_use_device_management(false);
model
} else {
// Use memory-efficient approach with blocks on CPU
info!("Using memory-efficient approach with blocks on CPU...");
let cpu_device = candle_core::Device::Cpu;
let cpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[flux_dev_or_schnell_path], dtype, &cpu_device)? };
let mut model = flux::model::Flux::new_with_gpu_core(&flux_config, cpu_vb, &device)?;
// Set active blocks
let active_blocks_count = active_blocks.unwrap_or(1);
model.set_max_active_blocks(active_blocks_count);
info!("Using {} active blocks on GPU", active_blocks_count);
// Enable prefetching if requested
if prefetch_next_batch.unwrap_or(false) {
model.set_prefetch_next_batch(true);
info!("Prefetching enabled - will overlap transfers with computation");
}
model
};
// Run denoising
flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, &timesteps, 4.)?
};
info!("FLUX denoising completed, unpacking latent image...");
let img_unpacked = flux::sampling::unpack(&img_latent, height, width)?;
// Load autoencoder for final decoding
info!("Loading autoencoder for final decoding...");
let ae_config = match flux_model {
FluxModel::Dev => flux::autoencoder::Config::dev(),
FluxModel::Schnell => flux::autoencoder::Config::schnell(),
};
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[flux_autoencoder_path], dtype, &device)? };
let ae_model = flux::autoencoder::AutoEncoder::new(&ae_config, vb)?;
info!("Decoding final image...");
let img = ae_model.decode(&img_unpacked)?;
// Convert tensor to image
info!("Converting tensor to image...");
let img = ((img.clamp(-1f32, 1f32)? + 1.0)? * 127.5)?.to_dtype(candle_core::DType::U8)?;
// Convert to RGB image
let image = tensor_to_image_buffer(&img.i(0)?)?;
info!("FLUX pipeline completed successfully");
Ok(image)
}
@@ -0,0 +1 @@
pub mod flux_pipeline;
+1
View File
@@ -1,4 +1,5 @@
pub mod background_removal;
pub mod flux;
pub mod image;
pub mod model_cache;
pub mod model_config;
@@ -20,12 +20,14 @@ use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
use crate::ml::stable_diffusion::infer_clip_text_embeddings::{infer_clip_text_embeddings, ClipArgs};
use crate::ml::stable_diffusion::remap_lcm_strength_range::remap_lcm_strength_range;
use image::DynamicImage;
use log::info;
use log::{debug, error, info};
use ml_weights_registry::weights_registry::weights::{LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
use rand::Rng;
const DEFAULT_STRENGTH : f64 = 75.0;
const DEFAULT_SEED : u64 = 42;
pub struct Args<'a, P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>> {
pub image: &'a DynamicImage,
pub prompt: String,
@@ -85,10 +87,9 @@ pub fn lcm_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<Rg
// Use LCM Scheduler instead of Euler Ancestral for better speed and quality
let mut scheduler = LCMScheduler::new(scheduler_steps, img2img_strength, LCMSchedulerConfig::default())?;
let seed = maybe_seed.unwrap_or_else(|| rand::thread_rng().gen());
//let seed = maybe_seed.unwrap_or_else(|| rand::thread_rng().gen());
let seed = maybe_seed.unwrap_or(DEFAULT_SEED);
let seed = 42;
device.set_seed(seed)?;
let guidance_scale = match cfg_scale {
@@ -129,13 +130,13 @@ pub fn lcm_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<Rg
tensor
};
info!("Text embeddings shape: {:?}", text_embeddings.shape());
debug!("Text embeddings shape: {:?}", text_embeddings.shape());
info!("Loading input image into tensor...");
debug!("Loading input image into tensor...");
let input_image = dynamic_image_to_tensor(image, &device, dtype)?;
info!("Reference image shape: {:?}", input_image.shape());
debug!("Reference image shape: {:?}", input_image.shape());
let vae_scale = get_vae_scale(sd_version);
@@ -161,7 +162,7 @@ pub fn lcm_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<Rg
let unet = match maybe_unet {
Some(unet) => unet,
None => {
info!("No unet found in cache; loading...");
info!("Building UNET model from file {:?}...", unet_path.as_ref());
let in_channels = match sd_version {
StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::V1_5Inpaint => 9,
@@ -217,16 +218,17 @@ pub fn lcm_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<Rg
let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &device, dtype)?;
// Generate latents from input image using the LCM approach
info!("Generating latents from input image...");
debug!("Generating latents from input image...");
let init_latents = (init_latent_dist.sample()? * vae_scale)?;
info!("Initial latents shape: {:?}", init_latents.shape());
debug!("Initial latents shape: {:?}", init_latents.shape());
// Calculate timesteps for LCM with proper img2img handling
let timesteps = LCMScheduler::get_timesteps_for_steps(scheduler_steps, img2img_strength);
let t_start = timesteps[0]; // First denoising step
// Add noise at the appropriate timestep for img2img
info!("Adding noise to latents for t_start={}", t_start);
debug!("Adding noise to latents for t_start={}", t_start);
let noise = init_latents.randn_like(0f64, 1f64)?;
let latents = scheduler.add_noise(&init_latents, noise, t_start)?;
@@ -244,7 +246,7 @@ pub fn lcm_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<Rg
let mut noise_pred = match unet.forward_with_guidance(&latent_model_input, timestep as f64, &text_embeddings, Some(&guidance_scale_embedding)) {
Ok(pred) => pred,
Err(e) => {
info!("UNet inference failed with error: {}", e);
error!("UNet inference failed with error: {}", e);
return Err(anyhow::anyhow!("UNet inference failed: {}", e));
},
};
@@ -259,21 +261,19 @@ pub fn lcm_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<Rg
latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, scheduler_steps)?;
}
info!("Diffusion process completed, decoding image...");
debug!("Diffusion process completed, decoding image...");
let image = vae.decode(&(latents / vae_scale)?)?;
info!("VAE decode completed");
info!("Scaling image back");
debug!("Scaling image back");
let image = ((image / 2.)? + 0.5)?;
info!("Normalized image values");
debug!("Normalized image values");
let image = (image.clamp(0f32, 1.)? * 255.)?;
info!("Convert to int8");
debug!("Convert to int8");
let image = image.to_dtype(DType::U8)?;
info!("Converted to 8-bit format");
debug!("Converted to 8-bit format");
let image = tensor_to_image_buffer(&image.i(0)?)?;
@@ -59,14 +59,6 @@ pub fn stable_diffusion_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>)
clip_weights_path,
} = args;
println!("Starting image generation with the following configuration:");
println!(" Model: {:?}", sd_version);
println!(" Prompt: {}", prompt);
println!(" Steps: {}", scheduler_steps);
println!(" Device: {:?}", device);
println!("Model dimensions: {}x{}", sd_config.width, sd_config.height);
// TODO(bt,2025-02-18): The scheduler is `EulerAncestralDiscreteScheduler`, but we may want to port an LCM scheduler.
// This is a target for performance improvement
// See: https://github.com/huggingface/candle/issues/1331
@@ -0,0 +1,80 @@
use crate::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction};
use crate::weights_registry::weight_descriptor_builder::weight;
/// https://huggingface.co/black-forest-labs/FLUX.1-dev
pub const FLUX_DEV : WeightDescriptor = weight!(
"Flux-1 Dev",
"flux1-dev.safetensors",
WeightFunction::ImageGeneration
);
// TODO(bt): I think the autoencoder for Dev and Schnell are the same.
// They have the same checksum unless I mixed up the files when downloading.
/// https://huggingface.co/black-forest-labs/FLUX.1-dev
pub const FLUX_DEV_AUTOENCODER : WeightDescriptor = weight!(
"Flux-1 Dev (Autoencoder)",
"flux1-dev.ae.safetensors",
WeightFunction::ImageGeneration
);
/// https://huggingface.co/black-forest-labs/FLUX.1-schnell
pub const FLUX_SCHNELL : WeightDescriptor = weight!(
"Flux-1 Schnell",
"flux1-schnell.safetensors",
WeightFunction::ImageGeneration
);
/// https://huggingface.co/black-forest-labs/FLUX.1-schnell
pub const FLUX_SCHNELL_AUTOENCODER : WeightDescriptor = weight!(
"Flux-1 Schnell (Autoencoder)",
"flux1-schnell.ae.safetensors",
WeightFunction::ImageGeneration
);
/// https://huggingface.co/lmz/candle-flux
pub const LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF : WeightDescriptor = weight!(
"LMZ Flux Schnell Quantized",
"lmz_candle_flux_flux1-schnell.gguf",
WeightFunction::ImageGeneration
);
/// https://huggingface.co/google/t5-v1_1-xxl/tree/main
/// NB: I couldn't find the safetensors file here
/// NB: This should live under "refs/pr/2"
pub const GOOGLE_T5_V1_1_XXL_MODEL : WeightDescriptor = weight!(
"Google T5 v1.1 (XXL)",
"google_t5_v1_1_xxl_model.safetensors",
WeightFunction::TextEncoder
);
/// https://huggingface.co/google/t5-v1_1-xxl/tree/main
/// NB: I couldn't find the safetensors file here
/// NB: This should live under "refs/pr/2"
pub const GOOGLE_T5_V1_1_XXL_CONFIG : WeightDescriptor = weight!(
"Google T5 v1.1 (XXL) Config",
"google_t5_v1_1_xxl_config.json",
WeightFunction::TextEncoder
);
/// https://huggingface.co/lmz/mt5-tokenizers
pub const LMZ_T5_TOKENIZER_JSON : WeightDescriptor = weight!(
"LMZ T5 Tokenizer",
"lmz_mt5_tokenizers_t5-v1_1-xxl.tokenizer.json",
WeightFunction::TextTokenizer
);
/// https://huggingface.co/openai/clip-vit-large-patch14/tree/main
pub const OPENAI_CLIP_VIT_P14_MODEL : WeightDescriptor = weight!(
"OpenAI CLIP VIT (patch 14) model",
"openai_clip-vit-large-patch14.model.safetensors",
WeightFunction::TextTokenizer
);
/// https://huggingface.co/openai/clip-vit-large-patch14/tree/main
pub const OPENAI_CLIP_VIT_P14_TOKENIZER_JSON: WeightDescriptor = weight!(
"OpenAI CLIP VIT (patch 14) tokenizer",
"openai_clip-vit-large-patch14.tokenizer.json",
WeightFunction::TextTokenizer
);
@@ -1,3 +1,4 @@
mod weight_descriptor_builder;
pub mod flux_weights;
pub mod weight_descriptor;
pub mod weights;
@@ -5,6 +5,7 @@
#[derive(Clone, Copy)]
pub enum WeightFunction {
ImageGeneration,
ImageSegmentation,
TextEncoder,
TextTokenizer,
@@ -29,6 +30,7 @@ pub struct WeightDescriptor {
impl WeightFunction {
pub fn as_descriptive_str(&self) -> &'static str {
match self {
WeightFunction::ImageGeneration => "image generation",
WeightFunction::ImageSegmentation => "image segmentation",
WeightFunction::TextEncoder => "text encoder",
WeightFunction::TextTokenizer => "text tokenizer",
+6 -1
View File
@@ -21,6 +21,7 @@ tauri-build = { version = "2.0.4", features = [] }
[dependencies]
# Internal
crockford = { path = "../../lib/data/crockford" }
filesys = { path = "../../lib/files/filesys" }
memory_caching = { path = "../../lib/caching/memory_caching" }
ml_weights_registry = { path = "../ml_weights_registry" }
@@ -29,12 +30,16 @@ primitives = { path = "../../lib/data/primitives" }
# Internal ML
# TODO: Only on Windows and Linux
#ml_models = { path = "../ml_models", optional = true }
# Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332
cudarc = { version = "0.13.5", features = ["std", "cublas", "cublaslt", "curand", "driver", "nvrtc", "f16", "cuda-version-from-build-system", "dynamic-linking"], default-features=false }
# External
#openssl = { version = "0.10", features = ["vendored"] }
# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls
anyhow = "1.0"
base64 = "0.22.1"
bytes = "1.10"
chrono = "0.4.40"
const_format = "0.2.34"
directories = "6.0.0"
dirs = "6.0.0"
@@ -43,8 +48,8 @@ half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_di
image = "0.25.5"
log = "0.4"
once_cell = "1.20.3"
ort = { version = "1.16.2", features = ["cuda", "load-dynamic"] }
rand = "0.8.5"
# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls
reqwest = { version = "0.12.12", features = ["stream"], default-features = false }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
@@ -0,0 +1,13 @@
use log::info;
#[tauri::command]
pub async fn inpaint_image(
image: &str,
mask: &str,
prompt: String,
) -> Result<String, String> {
info!("inpaint_endpoint endpoint called.");
Ok(mask.to_string())
}
@@ -1,5 +1,7 @@
pub mod download_models;
pub mod flip_image;
pub mod inpaint_image_endpoint;
pub mod realtime_image_endpoint;
pub mod remove_background_endpoint;
pub mod save_image_endpoint;
pub mod text_to_image_endpoint;
@@ -31,6 +31,7 @@ pub async fn infer_image(
image: &str,
prompt: Option<String>,
strength: Option<f64>,
seed: Option<u64>,
model_config: State<'_, AppConfig>,
model_cache: State<'_, ModelCache>,
prompt_cache: State<'_, PromptCache>,
@@ -49,12 +50,15 @@ pub async fn infer_image(
// TODO(bt,2025-02-17): Running out of vram with full image buffer size
let mut image = image.resize(512, 512, FilterType::CatmullRom);
let seed = seed.or(model_config.seed);
#[cfg(not(target_os = "macos"))]
{
let result = infer_image_impl(
&prompt,
image,
strength,
seed,
&model_config,
&model_cache,
prompt_cache,
@@ -83,6 +87,7 @@ async fn infer_image_impl(
prompt: &str,
image: DynamicImage,
strength: Option<f64>,
seed: Option<u64>,
config: &AppConfig,
model_cache: &ModelCache,
prompt_cache: State<'_, PromptCache>,
@@ -106,7 +111,7 @@ async fn infer_image_impl(
cfg_scale: config.cfg_scale,
use_flash_attn: true,
model_config: &config.model_config,
maybe_seed: None,
maybe_seed: seed,
scheduler_steps: config.scheduler_steps,
vae_path: &vae_path,
unet_path: &unet_path,
@@ -0,0 +1,38 @@
use crate::state::app_config::AppConfig;
use crate::state::app_dir::AppDataRoot;
use crate::stubs::model_cache::ModelCache;
use crate::stubs::prompt_cache::PromptCache;
use crate::utils::image::decode_base64_image::decode_base64_image;
use chrono::Local;
use crockford::crockford_entropy_lower;
use image::{ImageFormat, RgbImage};
use log::info;
use tauri::{AppHandle, State};
const FILENAME_LENGTH : usize = 10;
/// This handler takes an image (as a base64 encoded string) and saves it.
#[tauri::command]
pub async fn save_image(
image: &str,
app_data_root: State<'_, AppDataRoot>,
) -> Result<String, String> {
info!("save_image endpoint called.");
let image = decode_base64_image(image)
.map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
let assets_dir = app_data_root.assets_dir();
let date_save_dir = assets_dir.make_or_get_current_date_dir()
.map_err(|err| format!("Couldn't create assets save directory: {}", err))?;
let entropy = crockford_entropy_lower(FILENAME_LENGTH); // TODO: These are ugly. A monotonic counter would be better.
let filename = format!("{}.png", entropy);
let filename = date_save_dir.join(filename);
image.save_with_format(filename, ImageFormat::Png)
.map_err(|err| format!("Couldn't save image: {}", err))?;
Ok("".to_string())
}
@@ -1,23 +1,21 @@
#[cfg(not(target_os = "macos"))]
use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
use ml_models::ml::flux::flux_pipeline::{flux_pipeline, FluxModel, FluxPipelineArgs};
use crate::state::app_config::AppConfig;
use crate::state::app_dir::AppDataRoot;
use crate::state::app_dir::{AppDataRoot, AppWeightsDir};
use crate::stubs::model_cache::ModelCache;
use crate::stubs::prompt_cache::PromptCache;
use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png;
use image::imageops::FilterType;
use image::{DynamicImage, ImageReader, RgbImage};
use log::{error, info};
use ml_weights_registry::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
use ml_weights_registry::weights_registry::flux_weights::{FLUX_DEV, FLUX_SCHNELL, FLUX_SCHNELL_AUTOENCODER, GOOGLE_T5_V1_1_XXL_CONFIG, GOOGLE_T5_V1_1_XXL_MODEL, LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF, LMZ_T5_TOKENIZER_JSON, OPENAI_CLIP_VIT_P14_MODEL, OPENAI_CLIP_VIT_P14_TOKENIZER_JSON};
use once_cell::sync::Lazy;
use std::io::Cursor;
use tauri::{AppHandle, State};
const RANDOM_SEED: u32 = 42;
const STRENGTH: f64 = 1.0;
const STRENGTH: f64 = 100.0;
const PNG_BYTES : &[u8] = include_bytes!("../../binary_includes/1024.png");
// NB: We're just hacking a temporary "text to image" test endpoint leveraging the existing image to image modality.
@@ -35,8 +33,8 @@ pub async fn text_to_image(
model_config: State<'_, AppConfig>,
model_cache: State<'_, ModelCache>,
prompt_cache: State<'_, PromptCache>,
app_data_root: State<'_, AppDataRoot>,
app: AppHandle,
app_data_root: State<'_, AppDataRoot>,
) -> Result<String, String> {
info!("text_to_image endpoint called.");
@@ -51,8 +49,8 @@ pub async fn text_to_image(
&model_config,
&model_cache,
prompt_cache,
app_data_root.weights_dir(),
app,
&app_data_root
).await;
image = match result {
@@ -80,35 +78,36 @@ async fn text_to_image_impl(
config: &AppConfig,
model_cache: &ModelCache,
prompt_cache: State<'_, PromptCache>,
weights_dir: &AppWeightsDir,
app: AppHandle,
app_data_root: &AppDataRoot,
) -> Result<RgbImage, String> {
let weights_dir = app_data_root.weights_dir();
let vae_path = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE);
let unet_path = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
let clip_json_path = weights_dir.weight_path(&CLIP_JSON);
let clip_weights_path= weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER);
let args = Args {
let args = FluxPipelineArgs {
image: &image,
prompt: prompt.to_string(),
uncond_prompt: "".to_string(),
model_cache,
prompt_cache: &prompt_cache,
app: &app,
flux_model: FluxModel::Schnell,
active_blocks: None,
use_cuda_stream: Some(true),
prefetch_next_batch: Some(false),
use_full_gpu: Some(true),
img2img_strength: strength,
cfg_scale: config.cfg_scale,
use_flash_attn: true,
seed: None,
use_quantized_model: Some(true),
model_config: &config.model_config,
maybe_seed: None,
scheduler_steps: config.scheduler_steps,
vae_path: &vae_path,
unet_path: &unet_path,
clip_json_path: &clip_json_path,
clip_weights_path: &clip_weights_path,
flux_dev_or_schnell_path: &weights_dir.weight_path(&FLUX_SCHNELL),
flux_schnell_quantized_gguf: &weights_dir.weight_path(&LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF),
flux_autoencoder_path: &weights_dir.weight_path(&FLUX_SCHNELL_AUTOENCODER),
google_t5_model_path: &weights_dir.weight_path(&GOOGLE_T5_V1_1_XXL_MODEL),
google_t5_config_path: &weights_dir.weight_path(&GOOGLE_T5_V1_1_XXL_CONFIG),
lmz_t5_tokenizer_path: &weights_dir.weight_path(&LMZ_T5_TOKENIZER_JSON),
openai_clip_model_path: &weights_dir.weight_path(&OPENAI_CLIP_VIT_P14_MODEL),
openai_clip_tokenizer_json_path: &weights_dir.weight_path(&OPENAI_CLIP_VIT_P14_TOKENIZER_JSON),
};
let image = lcm_pipeline(args)
let image = flux_pipeline(args)
.map_err(|err| format!("failure to encode image: {:?}", err))?;
Ok(image)
@@ -8,6 +8,9 @@ pub enum NotificationModelType {
TextTokenizer,
Unet,
Vae,
Json,
ClipEncoder,
FluxModel,
}
#[derive(Clone, Serialize)]
@@ -26,5 +29,17 @@ pub enum NotificationEvent<'a> {
ModelDownloadComplete {
model_name: &'a str,
model_type: NotificationModelType,
}
},
ModelLoading {
model_name: String,
model_type: NotificationModelType,
},
ModelInference {
model_name: String,
model_type: NotificationModelType,
},
ModelInferenceComplete {
model_name: String,
model_type: NotificationModelType,
},
}
+9 -2
View File
@@ -8,8 +8,10 @@ pub mod utils;
use crate::endpoints::download_models::download_models;
use crate::endpoints::flip_image::flip_image;
use crate::endpoints::inpaint_image_endpoint::inpaint_image;
use crate::endpoints::realtime_image_endpoint::infer_image;
use crate::endpoints::remove_background_endpoint::remove_background;
use crate::endpoints::save_image_endpoint::save_image;
use crate::endpoints::text_to_image_endpoint::text_to_image;
use crate::state::app_config::AppConfig;
use crate::stubs::model_cache::ModelCache;
@@ -28,7 +30,6 @@ pub fn run() {
let config = AppConfig::init()
.expect("config should load");
let prompt_cache = PromptCache::with_capacity(8)
.expect("prompt cache should load");
@@ -44,7 +45,10 @@ pub fn run() {
tauri::Builder::default()
.plugin(tauri_plugin_log::Builder::new()
.level(log::LevelFilter::Info)
.targets(vec![Target::new(TargetKind::Stdout)])
.targets(vec![
Target::new(TargetKind::Stdout),
Target::new(TargetKind::LogDir { file_name: Some(app_data_root.log_file_name_str().to_string()) }),
])
.build())
.setup(|app| {
// TODO(bt): This is broken on windows
@@ -71,7 +75,10 @@ pub fn run() {
download_models,
flip_image,
infer_image,
inpaint_image,
remove_background,
save_image,
text_to_image,
text_to_image,
])
.run(tauri::generate_context!())
@@ -4,6 +4,7 @@ use anyhow::anyhow;
use directories::UserDirs;
use ml_weights_registry::weights_registry::weight_descriptor::WeightDescriptor;
use std::path::{Path, PathBuf};
use chrono::{DateTime, Local};
use tempdir::TempDir;
use tempfile::{Builder, NamedTempFile};
@@ -13,6 +14,9 @@ const WEIGHTS_SUBDIRECTORY : &str = "weights";
const TEMPORARY_SUBDIRECTORY : &str = "temp";
/// Note: Tauri appends ".log" to the end of the filename.
const LOG_FILE_NAME : &str = "application_debug";
/// The path to the application data directory, which includes "asset" and "weights" data.
#[derive(Clone)]
pub struct AppDataRoot {
@@ -20,6 +24,8 @@ pub struct AppDataRoot {
assets_dir: AppAssetsDir,
weights_dir: AppWeightsDir,
temp_dir: TemporaryDir,
log_file_name: PathBuf,
log_file_name_string: String,
}
#[derive(Clone)]
@@ -70,12 +76,19 @@ impl AppDataRoot {
let assets_dir = AppAssetsDir::create_existing(dir.join(ASSETS_SUBDIRECTORY))?;
let weights_dir = AppWeightsDir::create_existing(dir.join(WEIGHTS_SUBDIRECTORY))?;
let temp_dir = TemporaryDir::create_existing(dir.join(TEMPORARY_SUBDIRECTORY))?;
let log_file_name = dir.join(LOG_FILE_NAME);
let log_file_name_string = log_file_name
.to_str()
.ok_or(anyhow!("couldn't convert log path to str"))?
.to_string();
Ok(Self {
path: dir,
assets_dir,
weights_dir,
temp_dir,
log_file_name,
log_file_name_string,
})
}
@@ -94,6 +107,14 @@ impl AppDataRoot {
pub fn path(&self) -> &Path {
&self.path
}
pub fn log_file_name(&self) -> &Path {
&self.log_file_name
}
pub fn log_file_name_str(&self) -> &str {
&self.log_file_name_string
}
}
impl AppAssetsDir {
@@ -117,6 +138,18 @@ impl AppAssetsDir {
pub fn path(&self) -> &Path {
&self.path
}
/// Return the current date asset directory
/// If it doesn't exist, create it.
pub fn make_or_get_current_date_dir(&self) -> anyhow::Result<PathBuf> {
let format: DateTime<Local> = Local::now();
let date_directory_name = format.format("%Y-%d-%m").to_string();
let full_path = self.path.join(date_directory_name);
if !full_path.exists() {
std::fs::create_dir(&full_path)?;
}
Ok(full_path)
}
}
impl AppWeightsDir {
@@ -13,6 +13,7 @@ use std::sync::{Arc, Mutex};
use tauri::{AppHandle, Emitter};
use tempfile::NamedTempFile;
use tokio::task::JoinHandle;
use ml_weights_registry::weights_registry::flux_weights::{FLUX_SCHNELL, FLUX_SCHNELL_AUTOENCODER, GOOGLE_T5_V1_1_XXL_CONFIG, GOOGLE_T5_V1_1_XXL_MODEL, LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF, LMZ_T5_TOKENIZER_JSON, OPENAI_CLIP_VIT_P14_MODEL, OPENAI_CLIP_VIT_P14_TOKENIZER_JSON};
const NOTIFICATION_CHANNEL_NAME : &str = "notification";
const MAX_FILES : usize = 8;
@@ -50,6 +51,16 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> !
//download_queue.push_back(SDXL_TURBO_VAE); // TODO(bt): REMOVE
//download_queue.push_back(SDXL_TURBO_UNET); // TODO(bt): REMOVE
// Flux:
download_queue.push_back(FLUX_SCHNELL);
download_queue.push_back(FLUX_SCHNELL_AUTOENCODER);
download_queue.push_back(LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF);
download_queue.push_back(GOOGLE_T5_V1_1_XXL_MODEL);
download_queue.push_back(GOOGLE_T5_V1_1_XXL_CONFIG);
download_queue.push_back(LMZ_T5_TOKENIZER_JSON);
download_queue.push_back(OPENAI_CLIP_VIT_P14_MODEL);
download_queue.push_back(OPENAI_CLIP_VIT_P14_TOKENIZER_JSON);
let mut handles = FuturesUnordered::new();
while let Some(model) = download_queue.pop_front() {
@@ -82,6 +93,7 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> !
WeightFunction::Unet => NotificationModelType::Unet,
WeightFunction::Vae => NotificationModelType::Vae,
WeightFunction::ImageSegmentation => NotificationModelType::ImageSegmentation,
WeightFunction::ImageGeneration => NotificationModelType::ImageSegmentation, // TODO FIX
};
let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadStarted {
@@ -0,0 +1,13 @@
use actix_cors::Cors;
use crate::util::netlify_branch_domain_matches::netlify_branch_domain_matches;
pub fn add_artcraft_2d(cors: Cors, _is_production: bool) -> Cors {
cors
// Hypothetical domains
.allowed_origin("https://2d.storyteller.ai")
.allowed_origin("https://2d.getartcraft.com")
// Netlify project
.allowed_origin_fn(|origin, _req_head| {
netlify_branch_domain_matches(origin, "storyteller-2d.netlify.app")
})
}
@@ -0,0 +1,13 @@
use actix_cors::Cors;
use crate::util::netlify_branch_domain_matches::netlify_branch_domain_matches;
pub fn add_artcraft_3d(cors: Cors, _is_production: bool) -> Cors {
cors
// Hypothetical domains
.allowed_origin("https://3d.storyteller.ai")
.allowed_origin("https://3d.getartcraft.com")
// Netlify project
.allowed_origin_fn(|origin, _req_head| {
netlify_branch_domain_matches(origin, "storyteller-3d.netlify.app")
})
}
@@ -1,3 +1,5 @@
pub (crate) mod artcraft_2d;
pub (crate) mod artcraft_3d;
pub (crate) mod development_only;
pub (crate) mod fakeyou;
pub (crate) mod legacy;
+6 -1
View File
@@ -2,7 +2,8 @@ use actix_cors::Cors;
use log::info;
use reusable_types::server_environment::ServerEnvironment;
use crate::configs::artcraft_2d::add_artcraft_2d;
use crate::configs::artcraft_3d::add_artcraft_3d;
use crate::configs::development_only::add_development_only;
use crate::configs::fakeyou::{add_fakeyou, add_fakeyou_dev_proxy};
use crate::configs::legacy::{add_legacy_storyteller_stream, add_legacy_trumped, add_legacy_vocodes, add_power_stream};
@@ -40,6 +41,10 @@ fn do_build_cors_config(is_production: bool) -> Cors {
cors = add_storyteller_board(cors, is_production);
cors = add_storyteller_render(cors, is_production);
// Artcraft
cors = add_artcraft_2d(cors, is_production);
cors = add_artcraft_3d(cors, is_production);
// Legacy
cors = add_legacy_trumped(cors, is_production);
cors = add_power_stream(cors, is_production);
@@ -234,17 +234,31 @@ impl BucketClient {
.map(|s| s.to_string())
.ok_or(anyhow!("could not convert object path to string"))?;
info!("downloading from bucket: {:?}", &object_path_str);
info!("creating file for bucket download: {:?}", filesystem_path.as_ref());
let mut output_file = tokio::fs::File::create(filesystem_path).await?;
let mut output_file = File::create(filesystem_path).await?;
let status_code = self.bucket.get_object_to_writer(&object_path_str, &mut output_file).await?;
let result = self.bucket.get_object_to_writer(&object_path_str, &mut output_file).await;
info!("downloading from bucket (named '{}'), path: {}", &self.bucket.name, &object_path_str);
let status_code = match result {
Ok(status_code) => status_code,
Err(err) => {
return bail!("Error downloading from bucket (named '{}'): {:?}", &self.bucket.name, err)
}
};
match status_code {
404 => bail!("File not found in bucket: {}", &object_path_str),
_ => {},
404 => {
error!("File not found in bucket (named '{}'), path: {}", &self.bucket.name, &object_path_str);
bail!("File not found in bucket: {}", &object_path_str)
},
_ => {
info!("download code: {}", status_code);
},
}
info!("download code: {}", status_code);
Ok(())
}
}
@@ -0,0 +1,35 @@
[package]
name = "openai_sora_client"
edition = "2021"
version = "0.0.1"
publish = false
[package.metadata.cargo-udeps.ignore]
normal = ["my-workspace-hack"]
[lib]
name = "openai_sora_client"
path = "src/lib.rs"
[dependencies]
# Workspace hack (faster compile times)
my-workspace-hack = { workspace = true }
# Internal
crockford = { path = "../../data/crockford" }
errors = { path = "../../errors" }
subprocess_common = { path = "../../subprocess_common" }
# External
anyhow = { workspace = true }
log = { workspace = true }
reqwest = { workspace = true, features = ["multipart"] }
serde = { workspace = true }
serde_derive = { workspace = true }
serde_json = { workspace = true }
tokio = { version = "1.38.0", features = ["macros"] }
[dev-dependencies]
testing = { path = "../../testing" }
filesys = { path = "../../files/filesys" }
@@ -0,0 +1,43 @@
use reqwest::RequestBuilder;
use errors::AnyhowResult;
// NB: It appears that the sentinel may require a matching user agent.
pub (crate) const USER_AGENT : &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36";
#[derive(Clone)]
pub struct SoraCredentials {
pub bearer_token: String,
pub cookie: String,
pub sentinel: Option<String>,
}
impl SoraCredentials {
pub fn authorization_header_value(&self) -> String {
match self.bearer_token.get(0..6) {
Some("bearer") | Some("Bearer") => self.bearer_token.clone(),
_ => "Bearer ".to_owned() + &self.bearer_token,
}
}
pub fn from_env() -> AnyhowResult<Self> {
let bearer = std::env::var("SORA_BEARER_TOKEN")?;
let cookie = std::env::var("SORA_COOKIE")?;
let sentinel = std::env::var("SORA_SENTINEL").ok();
Ok(Self { bearer_token: bearer, cookie, sentinel })
}
pub fn add_credential_headers_to_request(&self, request: RequestBuilder) -> RequestBuilder {
let bearer_header = self.authorization_header_value();
let mut request = request
.header("User-Agent", USER_AGENT)
.header("Cookie", &self.cookie)
.header("Authorization", bearer_header);
if let Some(sentinel) = &self.sentinel {
request = request.header("OpenAI-Sentinel-Token", sentinel);
}
request
}
}
@@ -0,0 +1,48 @@
#[derive(Debug, Clone)]
pub struct SoraImageGenResponse {
pub task_id: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NumImages {
One,
Two,
Four
}
impl NumImages {
pub fn as_count(&self) -> usize {
match self {
NumImages::One => 1,
NumImages::Two => 2,
NumImages::Four => 4,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ImageSize {
Square,
Wide,
Tall,
}
impl ImageSize {
// TODO: Verify these dimensions. 480x480 is correct.
pub fn as_dimensions(&self) -> (u16, u16) {
match self {
ImageSize::Square => (480, 480),
ImageSize::Wide => (854, 480),
ImageSize::Tall => (480, 854),
}
}
pub fn as_width(&self) -> u16 {
self.as_dimensions().0
}
pub fn as_height(&self) -> u16 {
self.as_dimensions().1
}
}
@@ -0,0 +1,130 @@
use anyhow::anyhow;
use serde_derive::{Deserialize, Serialize};
use errors::AnyhowResult;
use crate::credentials::SoraCredentials;
const SORA_IMAGE_GEN_URL: &str = "https://sora.com/backend/video_gen";
/// This user agent is tied to the sentinel generation. If we need to change it, we may need to change sentinel generation too.
const USER_AGENT : &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36";
#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub (crate) enum VideoGenType {
ImageGen,
}
#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub (crate) enum OperationType {
/// Simple prompt without reference images
SimpleCompose,
/// Prompt with reference images
Remix,
}
#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub (crate) enum InpaintItemType {
Image,
}
#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub (crate) struct InpaintItem {
pub r#type: InpaintItemType,
/// Typically "0"
pub frame_index: u32,
/// Unknown field; defaults to null
pub preset_id: Option<String>,
/// Unknown field; defaults to null
pub generation_id: Option<String>,
/// Token identifier of the item, eg. "media_01jqt9vt20erx9zvryf3v1pecx"
pub upload_media_id: String,
/// Typically "0"
pub source_start_frame: u32,
/// Typically "0"
pub source_end_frame: u32,
/// Unknown field; defaults to null
pub crop_bounds: Option<String>,
}
#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub (crate) struct RawSoraImageGenRequest {
/// eg. "image_gen"
pub r#type: VideoGenType,
/// eg. "simple_compose"
pub operation: OperationType,
/// The user's raw prompt
pub prompt: String,
/// The number of variants to generate, eg. 2
pub n_variants: usize,
/// The width of the image, eg. 480
pub width: u16,
/// The height of the image, eg. 480
pub height: u16,
/// The number of frames to generate, eg. 1
pub n_frames: u8,
/// The items to inpaint, eg. []
pub inpaint_items: Vec<InpaintItem>,
}
#[derive(Deserialize)]
#[serde(rename_all = "snake_case")]
pub (crate) struct RawSoraResponse {
/// eg. "task_01jqsz9dsae9tvjygf1abrv3xf"
pub id: String,
/// not known
pub priority: Option<String>,
}
// https://sora.com/backend/notif?limit=100&before=task_01jqpg8qghenet0rw2a79p0vbn
// error
// message "Missing bearer authentication in header"
// type "invalid_request_error"
// param null
// code null
#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub (crate) struct RawSoraErrorResponse {
error: String,
message: String,
r#type: String,
param: Option<String>,
code: Option<String>,
}
/// Don't expose the internal request implementation as there are only a few "correct" ways to call the API.
pub (crate) async fn image_gen_http_request(sora_request: RawSoraImageGenRequest, credentials: &SoraCredentials) -> AnyhowResult<RawSoraResponse> {
let client = reqwest::Client::new();
let sentinel = match credentials.sentinel.as_deref() {
None => return Err(anyhow!("Calls to image gen require a sentinel in the credentials payload.")),
Some(sentinel) => sentinel,
};
let http_request = client.post(SORA_IMAGE_GEN_URL)
.header("User-Agent", USER_AGENT)
.header("Cookie", &credentials.cookie)
.header("Authorization", credentials.authorization_header_value())
.header("Content-Type", "application/json")
.header("OpenAI-Sentinel-Token", sentinel);
let http_request = http_request.json(&sora_request).build()?;
let response = client.execute(http_request)
.await?;
let status = response.status();
let response_body = &response.text().await?;
if status != reqwest::StatusCode::OK {
return Err(anyhow::anyhow!("the request failed; status = {:?}, message = {:?}", status, response_body));
}
let response = serde_json::from_str(response_body)?;
Ok(response)
}
@@ -0,0 +1,4 @@
pub (crate) mod image_gen_http_request;
pub mod common;
pub mod sora_image_gen_remix;
pub mod sora_image_gen_simple;
@@ -0,0 +1,87 @@
use errors::AnyhowResult;
use crate::credentials::SoraCredentials;
use crate::image_gen::common::{ImageSize, NumImages, SoraImageGenResponse};
use crate::image_gen::image_gen_http_request::{image_gen_http_request, InpaintItem, InpaintItemType, OperationType, RawSoraImageGenRequest, VideoGenType};
pub struct SoraImageGenRemixRequest<'a> {
pub prompt: String,
pub num_images: NumImages,
pub image_size: ImageSize,
pub sora_media_tokens: Vec<String>,
pub credentials: &'a SoraCredentials,
}
/// The "remix" commands let you supply additional images as context.
/// Sora "media tokens" of previously uploaded images must be supplied.
pub async fn sora_image_gen_remix(request: SoraImageGenRemixRequest<'_>) -> AnyhowResult<SoraImageGenResponse> {
let args = RawSoraImageGenRequest {
r#type: VideoGenType::ImageGen,
operation: OperationType::Remix,
prompt: request.prompt,
n_variants: request.num_images.as_count(),
width: request.image_size.as_width(),
height: request.image_size.as_height(),
n_frames: 1,
inpaint_items: request.sora_media_tokens.into_iter().map(|token| {
InpaintItem {
r#type: InpaintItemType::Image,
frame_index: 0,
preset_id: None,
generation_id: None,
upload_media_id: token,
source_start_frame: 0,
source_end_frame: 0,
crop_bounds: None,
}
}).collect(),
};
// TODO: Error handling.
let result = image_gen_http_request(args, &request.credentials).await?;
Ok(SoraImageGenResponse {
task_id: result.id,
})
}
#[cfg(test)]
mod tests {
use std::fs::read_to_string;
use errors::AnyhowResult;
use testing::test_file_path::test_file_path;
use crate::credentials::SoraCredentials;
use crate::image_gen::common::{ImageSize, NumImages};
use crate::image_gen::sora_image_gen_remix::{sora_image_gen_remix, SoraImageGenRemixRequest};
#[ignore] // You can manually run "ignore" tests in the IDE, but they won't run in CI.
#[tokio::test]
pub async fn manual_test() -> AnyhowResult<()> {
let sentinel = read_to_string(test_file_path("test_data/temp/sentinel.txt")?)?;
let sentinel = sentinel.trim().to_string();
let cookie = read_to_string(test_file_path("test_data/temp/cookie.txt")?)?;
let cookie = cookie.trim().to_string();
let bearer = read_to_string(test_file_path("test_data/temp/bearer.txt")?)?;
let bearer = bearer.trim().to_string();
let creds = SoraCredentials {
bearer_token: bearer,
cookie,
sentinel: Some(sentinel),
};
let response = sora_image_gen_remix(SoraImageGenRemixRequest {
prompt: "Match the pose and scene layout from the uploaded image exactly. A smart dog getting kisses from another dog. Brick building. Anime style".to_string(),
num_images: NumImages::One,
image_size: ImageSize::Square,
sora_media_tokens: vec!["media_01jqyhrz4detyvtzwp2p4j63ad".to_string()],
credentials: &creds,
}).await?;
println!("task_id: {}", response.task_id);
assert!(response.task_id.starts_with("task_"));
Ok(())
}
}
@@ -0,0 +1,72 @@
use errors::AnyhowResult;
use crate::credentials::SoraCredentials;
use crate::image_gen::common::{ImageSize, NumImages, SoraImageGenResponse};
use crate::image_gen::image_gen_http_request::{image_gen_http_request, OperationType, RawSoraImageGenRequest, VideoGenType};
pub struct SoraImageGenSimpleRequest<'a> {
pub prompt: String,
pub num_images: NumImages,
pub image_size: ImageSize,
pub credentials: &'a SoraCredentials,
}
pub async fn sora_image_gen_simple(args: SoraImageGenSimpleRequest<'_>) -> AnyhowResult<SoraImageGenResponse> {
let sora_request = RawSoraImageGenRequest {
r#type: VideoGenType::ImageGen,
operation: OperationType::SimpleCompose,
prompt: args.prompt,
n_variants: args.num_images.as_count(),
width: args.image_size.as_width(),
height: args.image_size.as_height(),
n_frames: 1,
inpaint_items: vec![],
};
// TODO: Error handling.
let result = image_gen_http_request(sora_request, args.credentials).await?;
Ok(SoraImageGenResponse {
task_id: result.id,
})
}
#[cfg(test)]
mod tests {
use std::fs::read_to_string;
use errors::AnyhowResult;
use testing::test_file_path::test_file_path;
use crate::credentials::SoraCredentials;
use crate::image_gen::common::{ImageSize, NumImages};
use crate::image_gen::sora_image_gen_simple::{sora_image_gen_simple, SoraImageGenSimpleRequest};
#[ignore] // You can manually run "ignore" tests in the IDE, but they won't run in CI.
#[tokio::test]
pub async fn manual_test() -> AnyhowResult<()> {
let sentinel = read_to_string(test_file_path("test_data/temp/sentinel.txt")?)?;
let sentinel = sentinel.trim().to_string();
let cookie = read_to_string(test_file_path("test_data/temp/cookie.txt")?)?;
let cookie = cookie.trim().to_string();
let bearer = read_to_string(test_file_path("test_data/temp/bearer.txt")?)?;
let bearer = bearer.trim().to_string();
let creds = SoraCredentials {
bearer_token: bearer,
cookie,
sentinel: Some(sentinel),
};
let response = sora_image_gen_simple(SoraImageGenSimpleRequest {
prompt: "A pirate and a ninja fight in a battle inside a UFO. Fully photo realistic, lifelike, lens flare".to_string(),
num_images: NumImages::One,
image_size: ImageSize::Square,
credentials: &creds,
}).await?;
println!("task_id: {}", response.task_id);
assert!(response.task_id.starts_with("task_"));
Ok(())
}
}
@@ -0,0 +1 @@
pub mod sora_job_status;
@@ -0,0 +1,64 @@
const USER_AGENT : &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36";
// https://sora.com/backend/notif?limit=100&before=task_01jqwwrkvgeqp8jsf5mqk1jceh
use crate::credentials::SoraCredentials;
pub async fn sora_job_status(
task_id: &str,
credentials: &SoraCredentials,
) -> anyhow::Result<()> {
let url = format!("https://sora.com/backend/notif?limit=100&before={}", task_id);
let client = reqwest::Client::new();
let request = client
.get(url)
.header("Authorization", &credentials.bearer_token)
.header("User-Agent", USER_AGENT)
.header("Cookie", &credentials.cookie);
let response = request.send().await?;
let json_response = &response.text().await?;
println!(" >>> response = {:?}", json_response);
//let response = serde_json::from_str(json_response)?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::fs::read_to_string;
use errors::AnyhowResult;
use testing::test_file_path::test_file_path;
use crate::credentials::SoraCredentials;
use crate::job::sora_job_status::sora_job_status;
#[ignore]
#[tokio::test]
pub async fn manual_test() -> AnyhowResult<()> {
let sentinel = read_to_string(test_file_path("test_data/temp/sentinel.txt")?)?;
let sentinel = sentinel.trim().to_string();
let cookie = read_to_string(test_file_path("test_data/temp/cookie.txt")?)?;
let cookie = cookie.trim().to_string();
let bearer = read_to_string(test_file_path("test_data/temp/bearer.txt")?)?;
let bearer = bearer.trim().to_string();
let creds = SoraCredentials {
bearer_token: bearer,
cookie,
sentinel: Some(sentinel),
};
let task_id = "task_01jqwwrkvgeqp8jsf5mqk1jceh";
let result = sora_job_status(task_id, &creds).await?;
Ok(())
}
}
@@ -0,0 +1,4 @@
pub mod credentials;
pub mod image_gen;
pub mod job;
pub mod upload;
@@ -0,0 +1,3 @@
pub (crate) mod upload_media_http_request;
pub mod upload_media_from_bytes;
pub mod upload_media_from_file;
@@ -0,0 +1,29 @@
use std::path::{Path, PathBuf};
use errors::AnyhowResult;
use crate::credentials::SoraCredentials;
use crate::upload::upload_media_http_request::{upload_media_http_request, SoraMediaUploadResponse};
pub struct SoraMediaUploadRequest<'a> {
pub file_path: String,
pub credentials: &'a SoraCredentials,
}
/// Upload bytes.
/// The underlying reqwest lib needs to own the bytes, so we can't pass as a reference.
pub async fn sora_media_upload_from_bytes<P: AsRef<Path>>(bytes: Vec<u8>, file_name: String, creds: &SoraCredentials) -> AnyhowResult<SoraMediaUploadResponse> {
let file_path = PathBuf::from(&file_name);
// TODO: Read file magic bytes first, then fall back to this.
let mime_type = match file_path.extension().and_then(|e| e.to_str()) {
Some("jpg") | Some("jpeg") => "image/jpeg",
Some("png") => "image/png",
// Some("webp") => "image/webp",
// Some("gif") => "image/gif",
// Some("mp4") => "video/mp4",
// Some("mov") => "video/quicktime",
// Some("webm") => "video/webm",
_ => "application/octet-stream",
};
Ok(upload_media_http_request(bytes, file_name, mime_type, creds).await?)
}
@@ -0,0 +1,81 @@
use std::path::Path;
use tokio::fs::File;
use tokio::io::AsyncReadExt;
use errors::AnyhowResult;
use crate::credentials::SoraCredentials;
use crate::upload::upload_media_http_request::{upload_media_http_request, SoraMediaUploadResponse};
/// Try to prevent buffer reallocations.
/// There's a better way to implement this.
const INITIAL_BUFFER_SIZE : usize = 1024*1024;
pub struct SoraMediaUploadRequest<'a> {
pub file_path: String,
pub credentials: &'a SoraCredentials,
}
pub async fn sora_media_upload_from_file<P: AsRef<Path>>(file_path: P, creds: &SoraCredentials) -> AnyhowResult<SoraMediaUploadResponse> {
let mut file = File::open(&file_path).await?;
let mut buffer = Vec::with_capacity(INITIAL_BUFFER_SIZE);
file.read_to_end(&mut buffer).await?;
let filename = file_path.as_ref().file_name()
.ok_or_else(|| anyhow::anyhow!("Invalid file path"))?
.to_string_lossy()
.to_string();
// TODO: Read file magic bytes first, then fall back to this.
let mime_type = match file_path.as_ref().extension().and_then(|e| e.to_str()) {
Some("jpg") | Some("jpeg") => "image/jpeg",
Some("png") => "image/png",
// Some("webp") => "image/webp",
// Some("gif") => "image/gif",
// Some("mp4") => "video/mp4",
// Some("mov") => "video/quicktime",
// Some("webm") => "video/webm",
_ => "application/octet-stream",
};
Ok(upload_media_http_request(buffer, filename, mime_type, creds).await?)
}
#[cfg(test)]
mod tests {
use std::fs::read_to_string;
use errors::AnyhowResult;
use testing::test_file_path::test_file_path;
use crate::credentials::SoraCredentials;
use crate::upload::upload_media_from_file::sora_media_upload_from_file;
#[ignore] // You can manually run "ignore" tests in the IDE, but they won't run in CI.
#[tokio::test]
pub async fn manual_test() -> AnyhowResult<()> {
let cookie = read_to_string(test_file_path("test_data/temp/cookie.txt")?)?;
let cookie = cookie.trim().to_string();
let bearer = read_to_string(test_file_path("test_data/temp/bearer.txt")?)?;
let bearer = bearer.trim().to_string();
let image_path = test_file_path("test_data/image/juno.jpg")?; // media_01jqyqgqpwf40tkcapq5bmaz5d
let creds = SoraCredentials {
bearer_token: bearer,
cookie,
sentinel: None,
};
let response = sora_media_upload_from_file(
image_path,
&creds,
).await?;
println!("media: {:?}", response);
println!("media_id: {}", response.id);
println!("media_url: {}", response.url);
assert!(response.id.starts_with("media_"));
Ok(())
}
}
@@ -0,0 +1,81 @@
use crate::credentials::{SoraCredentials, USER_AGENT};
use errors::AnyhowResult;
use reqwest::multipart::{Form, Part};
use reqwest::Client;
use serde::Deserialize;
const SORA_UPLOAD_MEDIA_URL: &str = "https://sora.com/backend/uploads";
// Content-Disposition: form-data; name="file"; filename="imagename.webp"
/*
Response:
{
"id": "media_01jqt9vt20erx9zvryf3v1pecx",
"type": "image",
"created_at": "2025-04-02T03:49:04.185867Z",
"filename": "Rith_essa.webp",
"extension": "jpeg",
"mime_type": "image/jpeg",
"url": "https://videos.openai.com/vg-assets/assets%2Fclient_upload%2Fmedia%2Fuser-9wf6JFdgRSJLjvSJ53LNAbV8%2Fmedia_01jqt9vt20erx9zvryf3v1pecx.jpeg?st=2025-04-02T02%3A39%3A01Z&se=2025-04-08T03%3A39%3A01Z&sks=b&skt=2025-04-02T02%3A39%3A01Z&ske=2025-04-08T03%3A39%3A01Z&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skoid=3d249c53-07fa-4ba4-9b65-0bf8eb4ea46a&skv=2019-02-02&sv=2018-11-09&sr=b&sp=r&spr=https%2Chttp&sig=%2BdnTQyw9OFF2kaVoEReT2cVfTqmKjGt18m7P4wKCWCI%3D&az=oaivgprodscus",
"width": 1200,
"height": 675,
"duration_sec": null,
"n_frames": 1,
"size_bytes": 118501,
"thumbnail_url": "https://videos.openai.com/vg-assets/assets%2Fclient_upload%2Fmedia%2Fuser-9wf6JFdgRSJLjvSJ53LNAbV8%2Fmedia_01jqt9vt20erx9zvryf3v1pecx.jpg?st=2025-04-02T02%3A39%3A01Z&se=2025-04-08T03%3A39%3A01Z&sks=b&skt=2025-04-02T02%3A39%3A01Z&ske=2025-04-08T03%3A39%3A01Z&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skoid=3d249c53-07fa-4ba4-9b65-0bf8eb4ea46a&skv=2019-02-02&sv=2018-11-09&sr=b&sp=r&spr=https%2Chttp&sig=9aU9xKI0BkhRtrTb0RD05%2BT6prBnmXk1I7nV2nTsFbY%3D&az=oaivgprodscus"
}
*/
#[derive(Debug, Deserialize)]
pub struct SoraMediaUploadResponse {
/// This is the media token id.
pub id: String,
pub r#type: String,
pub created_at: String,
pub filename: String,
pub extension: String,
pub mime_type: String,
pub url: String,
pub width: Option<u32>,
pub height: Option<u32>,
pub duration_sec: Option<f64>,
pub n_frames: Option<u32>,
pub size_bytes: u64,
pub thumbnail_url: Option<String>,
}
pub (crate) struct SoraMediaUploadRequest<'a> {
pub file_path: String,
pub credentials: &'a SoraCredentials,
}
pub (crate) async fn upload_media_http_request(file_bytes: Vec<u8>, filename: String, mime_type: &str, credentials: &SoraCredentials) -> AnyhowResult<SoraMediaUploadResponse> {
// Create multipart form
let part = Part::bytes(file_bytes) // NB: Reqwest needs to own the bytes.
.file_name(filename) // NB: Reqwest needs to own the bytes
.mime_str(mime_type)?;
let form = Form::new().part("file", part);
// Make API request
let client = Client::new();
let request_builder = client.post(SORA_UPLOAD_MEDIA_URL)
.multipart(form)
.header("User-Agent", USER_AGENT)
.header("Cookie", &credentials.cookie)
.header("Authorization", credentials.authorization_header_value());
let response = request_builder.send().await?;
// Check response status
if !response.status().is_success() {
let status = response.status();
let text = response.text().await?;
return Err(anyhow::anyhow!("Upload failed with status {}: {}", status, text));
}
// Parse response
let upload_response = response.json::<SoraMediaUploadResponse>().await?;
Ok(upload_response)
}
+30
View File
@@ -0,0 +1,30 @@
[package]
name = "image_algorithms"
edition = "2018"
version = "0.0.1"
authors = [
"Brandon Thomas <bt@brand.io>",
"Brandon Thomas <echelon@gmail.com>",
]
publish = false
[lib]
name = "image_algorithms"
path = "src/lib.rs"
[package.metadata.cargo-udeps.ignore]
normal = ["my-workspace-hack"]
[dependencies]
# Workspace hack (faster compile times)
my-workspace-hack = { version = "0.1", path = "../../my-workspace-hack" }
# Internal
errors = { path = "../errors" }
# External
image = "0.25.6"
[dev-dependencies]
testing = { path = "../testing" }
+1
View File
@@ -0,0 +1 @@
pub mod remove_image_border;
@@ -0,0 +1,78 @@
use image::{DynamicImage, Pixel};
use errors::AnyhowResult;
pub fn remove_image_border(image: &DynamicImage) -> AnyhowResult<DynamicImage> {
let mut image = image.to_rgba8();
let mut top = 0u32;
let mut bottom = image.height();
let mut left = 0u32;
let mut right = image.width();
// TODO: Iterate with binary search - it'll probably be much faster.
for (i, mut row) in image.rows().enumerate() {
let count_non_transparent = row.filter(|pixel| pixel.0[3] != 0).count();
if count_non_transparent < 5 {
top = i as u32;
} else {
break;
}
}
for (i, mut row) in image.rows().enumerate().rev() {
//let all_transparent = row.all(|pixel| pixel.0[3] == 0);
let count_non_transparent = row.filter(|pixel| pixel.0[3] != 0).count();
//if all_transparent {
if count_non_transparent < 5 {
bottom = i as u32;
} else {
break;
}
}
for (x, y, pixel) in image.enumerate_pixels() {
}
for (y, mut row) in image.enumerate_rows() {
for (x, y, pixel) in row {
}
}
let image = DynamicImage::ImageRgba8(image);
let width = image.width();
let height = bottom - top;
println!("top: {}, bottom: {}", top, bottom);
println!("orig width: {}, height: {}", image.width(), image.height());
println!("new width: {}, height: {}", width, height);
let image = image.crop_imm(0, top, width, height);
Ok(image)
}
#[cfg(test)]
mod tests {
use image::{ImageFormat, ImageReader};
use errors::AnyhowResult;
use testing::test_file_path::test_file_path;
use crate::remove_image_border::remove_image_border;
#[ignore] // NB: Not going to check these files into the repo. Just for manual testing.
#[test]
fn test_crop_image() -> AnyhowResult<()> {
let filename = test_file_path("test_data/image/mario.png")?;
let image = ImageReader::open(filename)?.decode()?;
let image = remove_image_border(&image)?;
image.save_with_format("testing_output.png", ImageFormat::Png)?;
Ok(())
}
}
@@ -1,7 +1,6 @@
use crate::payloads::generic_inference_args::inner_payloads::f5_tts_payload::F5TTSPayload;
use crate::payloads::generic_inference_args::inner_payloads::face_fusion_payload::FaceFusionPayload;
use crate::payloads::generic_inference_args::inner_payloads::gptsovits_payload::GptSovitsPayload;
use crate::payloads::generic_inference_args::inner_payloads::image_generation_payload::StableDiffusionArgs;
use crate::payloads::generic_inference_args::inner_payloads::lipsync_payload::LipsyncArgs;
use crate::payloads::generic_inference_args::inner_payloads::live_portrait_payload::LivePortraitPayload;
use crate::payloads::generic_inference_args::inner_payloads::mocap_payload::MocapArgs;
@@ -12,6 +11,8 @@ use crate::payloads::generic_inference_args::inner_payloads::tts_payload::TTSArg
use crate::payloads::generic_inference_args::inner_payloads::videofilter_payload::RerenderArgs;
use crate::payloads::generic_inference_args::inner_payloads::workflow_payload::WorkflowArgs;
use errors::AnyhowResult;
use crate::payloads::generic_inference_args::inner_payloads::image_generation_payload::StableDiffusionArgs;
use crate::payloads::generic_inference_args::inner_payloads::sora_image_gen_args::SoraImageGenArgs;
/// Used to encode extra state for the `generic_inference_jobs` table.
/// This should act somewhat like a serialized protobuf stored inside a record.
@@ -111,7 +112,7 @@ pub enum PolymorphicInferenceArgs {
Tts(TTSArgs),
// No arguments yet.
// It might be best to just not include this when not used.
/// Voice conversion. (Short name to save space when serializing.)
Vc {
/// Argument for so-vits-svc
@@ -143,6 +144,10 @@ pub enum PolymorphicInferenceArgs {
/// Image generation. (Short name to save space when serializing.)
Ig(StableDiffusionArgs),
/// Sora Image Generation. (Short name to save space when serializing.)
Sg(SoraImageGenArgs),
/// Mocap (Short name to save space when serializing.)
Mc(MocapArgs),
@@ -1,3 +1,4 @@
use tokens::tokens::media_files::MediaFileToken;
use tokens::tokens::model_weights::ModelWeightToken;
/// Video sources can be one of several:
@@ -52,7 +53,7 @@ pub struct StableDiffusionArgs {
#[serde(rename = "se")]
#[serde(skip_serializing_if = "Option::is_none")]
pub maybe_seed: Option<i64>,
#[serde(rename = "mu")]
#[serde(skip_serializing_if = "Option::is_none")]
pub maybe_upload_path: Option<String>,
@@ -60,7 +61,7 @@ pub struct StableDiffusionArgs {
#[serde(rename = "cf")]
#[serde(skip_serializing_if = "Option::is_none")]
pub maybe_cfg_scale: Option<u32>,
#[serde(rename = "lu")]
#[serde(skip_serializing_if = "Option::is_none")]
pub maybe_lora_upload_path: Option<String>,
@@ -1,3 +1,4 @@
pub mod f5_tts_payload;
pub mod face_fusion_payload;
pub mod gptsovits_payload;
pub mod image_generation_payload;
@@ -5,9 +6,9 @@ pub mod lipsync_payload;
pub mod live_portrait_payload;
pub mod mocap_payload;
pub mod render_engine_scene_to_video_payload;
pub mod seed_vc_payload;
pub mod sora_image_gen_args;
pub mod studio_gen2_payload;
pub mod tts_payload;
pub mod videofilter_payload;
pub mod workflow_payload;
pub mod f5_tts_payload;
pub mod seed_vc_payload;
@@ -0,0 +1,23 @@
use tokens::tokens::media_files::MediaFileToken;
#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct SoraImageGenArgs {
#[serde(rename = "p")]
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt: Option<String>,
#[serde(rename = "sm")]
pub scene_snapshot_media_token: Option<MediaFileToken>,
#[serde(rename = "mf")]
#[serde(skip_serializing_if = "Option::is_none")]
pub maybe_additional_media_file_tokens: Option<Vec<MediaFileToken>>,
#[serde(rename = "sa")]
#[serde(skip_serializing_if = "Option::is_none")]
pub maybe_number_of_samples: Option<u32>,
#[serde(rename = "su")]
#[serde(skip_serializing_if = "Option::is_none")]
pub maybe_sora_media_upload_tokens: Option<Vec<String>>,
}
+2
View File
@@ -17,7 +17,9 @@ pub fn test_file_path<P: AsRef<Path>>(path_from_repo_root: P) -> AnyhowResult<Pa
.to_str()
.ok_or(anyhow!("path cannot convert to string"))?;
path.push(format!("../../../{}", file));
println!(" >>> path = {:?}", path);
Ok(path.canonicalize()?)
}
#[cfg(test)]
@@ -69,6 +69,7 @@ tts_common = { path = "../../../lib/tts_common" }
url_config = { path = "../../../lib/url_config" }
user_input_common = { path = "../../../lib/user_input_common" }
videos = { path = "../../../lib/files/videos" }
openai_sora_client = { path = "../../../lib/api_clients/openai_sora_client" }
# Internal Components
billing_component = { path = "../../plugins/billing" }
@@ -124,6 +124,8 @@ use crate::http_server::endpoints::stats::get_unified_queue_stats_handler::*;
use crate::http_server::endpoints::tags::list_tags_for_entity_handler::*;
use crate::http_server::endpoints::tags::set_tags_for_entity_handler::*;
use crate::http_server::endpoints::tts::enqueue_infer_tts_handler::enqueue_infer_tts_handler::*;
use crate::http_server::endpoints::image_studio::prompt::enqueue_studio_image_generation_handler::*;
use crate::http_server::endpoints::image_studio::upload::upload_snapshot_media_file_handler::*;
use crate::http_server::endpoints::user_bookmarks::batch_get_user_bookmarks_handler::*;
use crate::http_server::endpoints::user_bookmarks::create_user_bookmark_handler::*;
use crate::http_server::endpoints::user_bookmarks::delete_user_bookmark_handler::*;
@@ -207,6 +209,8 @@ use crate::http_server::web_utils::response_success_helpers::*;
crate::http_server::endpoints::media_files::upload::upload_engine_asset::upload_engine_asset_media_file_handler::upload_engine_asset_media_file_handler,
crate::http_server::endpoints::media_files::upload::upload_generic::upload_media_file_handler::upload_media_file_handler,
crate::http_server::endpoints::media_files::upload::upload_image_media_file_handler::upload_image_media_file_handler,
crate::http_server::endpoints::image_studio::prompt::enqueue_studio_image_generation_handler::enqueue_studio_image_generation_handler,
crate::http_server::endpoints::image_studio::upload::upload_snapshot_media_file_handler::upload_snapshot_media_file_handler,
crate::http_server::endpoints::studio_gen2::enqueue_studio_gen2_handler::enqueue_studio_gen2_handler,
crate::http_server::endpoints::media_files::upload::upload_new_engine_asset_media_file_handler::upload_new_engine_asset_media_file_handler,
crate::http_server::endpoints::media_files::upload::upload_new_scene_media_file_handler::upload_new_scene_media_file_handler,
@@ -336,6 +340,9 @@ use crate::http_server::web_utils::response_success_helpers::*;
AppStateServerInfo,
AppStateStatusAlertCategory,
AppStateStatusAlertInfo,
EnqueueStudioImageGenRequest,
EnqueueImageGenRequestSuccessResponse,
EnqueueImageGenRequestError,
AppStateSubscriptionProductKey,
AppStateUserInfo,
AppStateUserLocale,
@@ -347,6 +354,8 @@ use crate::http_server::web_utils::response_success_helpers::*;
BatchGetMediaFilesQueryParams,
BatchGetMediaFilesSuccessResponse,
BatchGetUserBookmarksError,
UploadSnapshotMediaFileForm,
UploadSnapshotMediaFileSuccessResponse,
BatchGetUserBookmarksQueryParams,
BatchGetUserBookmarksResponse,
BatchGetUserRatingError,
@@ -0,0 +1 @@
pub mod set_sora_secret_handler;
@@ -0,0 +1,137 @@
use std::fmt;
use std::sync::Arc;
use actix_web::{HttpRequest, HttpResponse, web, Either};
use actix_web::error::ResponseError;
use actix_web::http::StatusCode;
use actix_web::web::{Form, Json, Path};
use anyhow::anyhow;
use log::{error, log, warn};
use r2d2_redis::redis::Commands;
use utoipa::ToSchema;
use enums::by_table::media_files::media_file_animation_type::MediaFileAnimationType;
use enums::by_table::media_files::media_file_engine_category::MediaFileEngineCategory;
use http_server_common::response::serialize_as_json_error::serialize_as_json_error;
use mysql_queries::queries::media_files::edit::update_media_file_animation_type::update_media_file_animation_type;
use mysql_queries::queries::media_files::get::get_media_file::get_media_file;
use tokens::tokens::media_files::MediaFileToken;
use crate::state::server_state::ServerState;
use crate::util::redis_sora_secrets::get_sora_credentials;
use crate::util::redis_sora_secrets::set_sora_credential_field;
use crate::util::redis_sora_secrets::RedisSoraCredentialSubkey;
#[derive(Deserialize, Clone, Copy)]
#[serde(rename_all = "snake_case")]
pub enum SetSoraSecretType {
Bearer,
Cookie,
Sentinel
}
#[derive(Deserialize, ToSchema)]
pub struct SetSoraSecretRequest {
pub key: SetSoraSecretType,
pub value: String,
}
#[derive(Serialize, ToSchema)]
pub struct SetSoraSecretSuccessResponse {
pub success: bool,
pub bearer: String,
pub cookie: String,
pub sentinel: String,
}
// =============== Error Response ===============
#[derive(Debug, Serialize, ToSchema)]
pub enum SetSoraSecretError {
BadInput(String),
NotAuthorized,
ServerError,
}
impl ResponseError for SetSoraSecretError {
fn status_code(&self) -> StatusCode {
match *self {
SetSoraSecretError::BadInput(_) => StatusCode::BAD_REQUEST,
SetSoraSecretError::NotAuthorized => StatusCode::UNAUTHORIZED,
SetSoraSecretError::ServerError => StatusCode::INTERNAL_SERVER_ERROR,
}
}
fn error_response(&self) -> HttpResponse {
serialize_as_json_error(self)
}
}
// NB: Not using derive_more::Display since Clion doesn't understand it.
impl fmt::Display for SetSoraSecretError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{:?}", self)
}
}
// =============== Handler ===============
/// Set a Sora secret
#[utoipa::path(
post,
tag = "Control Plane",
path = "/v1/control_plane/sora_secret",
responses(
(status = 200, description = "Success", body = SetSoraSecretSuccessResponse),
(status = 400, description = "Bad input", body = SetSoraSecretError),
(status = 401, description = "Not authorized", body = SetSoraSecretError),
(status = 500, description = "Server error", body = SetSoraSecretError),
),
params(
("request" = SetSoraSecretRequest, description = "Payload for Request"),
("path" = MediaFileTokenPathInfo, description = "Path for Request")
)
)]
pub async fn set_sora_secret_handler(
http_request: HttpRequest,
request: Form<SetSoraSecretRequest>,
server_state: web::Data<Arc<ServerState>>
) -> Result<Json<SetSoraSecretSuccessResponse>, SetSoraSecretError> {
// TODO(bt,2025-04-03): Some form of authentication should be required for this endpoint.
let value = request.value.trim().to_string();
let mut redis = server_state.redis_pool
.get()
.map_err(|e| {
error!("redis error: {:?}", e);
SetSoraSecretError::ServerError
})?;
set_sora_credential_field(
&mut redis,
match request.key {
SetSoraSecretType::Bearer => RedisSoraCredentialSubkey::Bearer,
SetSoraSecretType::Cookie => RedisSoraCredentialSubkey::Cookie,
SetSoraSecretType::Sentinel => RedisSoraCredentialSubkey::Sentinel,
},
&value
).map_err(|e| {
error!("redis error: {:?}", e);
SetSoraSecretError::ServerError
})?;
let credentials = get_sora_credentials(&mut redis)
.map_err(|e| {
error!("redis error: {:?}", e);
SetSoraSecretError::ServerError
})?;
// TODO(bt,2025-04-03): Probably not smart to return, but useful for debugging.
Ok(Json(SetSoraSecretSuccessResponse {
success: true,
bearer: credentials.bearer_token,
cookie: credentials.cookie,
sentinel: credentials.sentinel.unwrap_or_else(|| "".to_string()),
}))
}
@@ -507,6 +507,9 @@ fn validate_request(
Ok(())
}
fn inference_mode_from_http_request(http_request: &HttpRequest) -> Option<TypeOfInference> {
inference_mode_from_url_path(http_request.path())
}
@@ -0,0 +1,68 @@
use actix_web::HttpRequest;
use anyhow::anyhow;
use log::{error, info};
use r2d2_redis::r2d2::PooledConnection;
use r2d2_redis::RedisConnectionManager;
use errors::AnyhowResult;
use http_server_common::request::get_request_header_optional::get_request_header_optional;
use openai_sora_client::credentials::SoraCredentials;
use crate::util::redis_sora_secrets::get_sora_credentials;
const SORA_BEARER_HEADER_NAME : &str = "sora-bearer";
const SORA_COOKIE_HEADER_NAME : &str = "sora-cookie";
const SORA_SENTINEL_HEADER_NAME : &str = "sora-sentinel";
pub fn get_sora_credentials_from_request(
http_request: &HttpRequest,
redis: &mut PooledConnection<RedisConnectionManager>,
) -> AnyhowResult<SoraCredentials> {
let sora_bearer = get_request_header_optional(&http_request, SORA_BEARER_HEADER_NAME);
let sora_cookie = get_request_header_optional(&http_request, SORA_COOKIE_HEADER_NAME);
let sora_sentinel = get_request_header_optional(&http_request, SORA_SENTINEL_HEADER_NAME);
match (&sora_bearer, &sora_cookie, &sora_sentinel) {
(Some(bearer), Some(cookie), Some(sentinel)) => {
info!("Sora credentials were present in HTTP request headers!");
return Ok(SoraCredentials {
bearer_token: bearer.to_string(),
cookie: cookie.to_string(),
sentinel: Some(sentinel.to_string()),
});
}
_ => {} // Fall through
}
let redis_result = get_sora_credentials(redis);
match redis_result {
Ok(mut credentials) => {
if let Some(bearer) = sora_bearer {
info!("override bearer value");
credentials.bearer_token = bearer;
}
if let Some(cookie) = sora_cookie {
info!("override cookie value");
credentials.cookie = cookie;
}
if let Some(sentinel) = sora_sentinel {
info!("override sentinel value");
credentials.sentinel = Some(sentinel);
}
return Ok(credentials)
},
Err(err) => {
error!("Failed to get Sora credentials from Redis; using env as backup. Error: {:?}", err);
// NB: Fall through
}
};
match SoraCredentials::from_env() {
Ok(creds) => Ok(creds),
Err(err) => {
error!("Failed to load Sora credentials from environment: {:?}", err);
Err(anyhow!("Failed to load Sora credentials from environment"))
}
}
}
@@ -0,0 +1,3 @@
pub mod upload;
pub mod prompt;
pub mod get_sora_credentials_from_request;
@@ -0,0 +1,465 @@
#![forbid(unused_mut)]
use std::fmt::Debug;
use std::fmt::{Display, Formatter};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use actix_web::error::ResponseError;
use actix_web::http::StatusCode;
use actix_web::web::Json;
use actix_web::{web, HttpRequest, HttpResponse};
use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
use cloud_storage::bucket_client::BucketClient;
use log::{debug, error, info};
use log::warn;
use mysql_queries::queries::media_files::get::get_media_file::get_media_file;
use openai_sora_client::credentials::SoraCredentials;
use serde::Deserialize;
use serde::Serialize;
use sqlx::MySqlPool;
use tempdir::TempDir;
use utoipa::ToSchema;
use web::Data;
use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType;
use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType;
use enums::common::visibility::Visibility;
use http_server_common::request::get_request_header_optional::get_request_header_optional;
use http_server_common::request::get_request_ip::get_request_ip;
use mysql_queries::payloads::generic_inference_args::generic_inference_args::{
GenericInferenceArgs,
InferenceCategoryAbbreviated,
PolymorphicInferenceArgs,
};
use mysql_queries::payloads::generic_inference_args::inner_payloads::sora_image_gen_args::SoraImageGenArgs;
use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{
insert_generic_inference_job,
InsertGenericInferenceArgs,
};
use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
use openai_sora_client::image_gen::common::{ImageSize, NumImages};
use openai_sora_client::image_gen::sora_image_gen_remix::{sora_image_gen_remix, SoraImageGenRemixRequest};
use openai_sora_client::upload::upload_media_from_file::{sora_media_upload_from_file, SoraMediaUploadRequest};
use tokens::tokens::generic_inference_jobs::InferenceJobToken;
use tokens::tokens::media_files::MediaFileToken;
use tokens::tokens::model_weights::ModelWeightToken;
use tokens::tokens::users::UserToken;
use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session;
use crate::http_server::endpoints::image_studio::get_sora_credentials_from_request::get_sora_credentials_from_request;
use crate::http_server::endpoints::tts::get_tts_inference_job_status::GetTtsInferenceStatusError;
use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format;
use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
use crate::state::server_state::ServerState;
use crate::util::allowed_studio_access::allowed_studio_access;
use crate::util::redis_sora_secrets::get_sora_credentials;
/// This is the number of images (batch size) to generate for each request.
/// We should allow all users to have multiple images generated at once as this
/// is what other providers do.
const MINIMUM_IMAGE_COUNT : u32 = 1;
/// The maximum number of images (batch size) to generate for each request.
const MAXIMUM_IMAGE_COUNT : u32 = 2;
/// Debug requests can get routed to special "debug-only" workers, which can
/// be used to trial new code, run debugging, etc.
const DEBUG_HEADER_NAME: &str = "enable-debug-mode";
/// The routing tag header can send workloads to particular k8s hosts.
/// This is useful for catching the live logs or intercepting the job.
const ROUTING_TAG_HEADER_NAME: &str = "routing-tag";
#[derive(Deserialize, ToSchema)]
pub struct EnqueueStudioImageGenRequest {
pub uuid_idempotency_token: String,
/// Image media file; the engine or canvas snapshot (screenshot).
pub snapshot_media_token: MediaFileToken,
/// The user's image generation prompt.
pub prompt: String,
/// Additional images to include (optional). Up to nine images.
pub maybe_additional_images: Option<Vec<MediaFileToken>>,
pub maybe_number_of_samples: Option<u32>,
}
#[derive(Serialize, ToSchema)]
pub struct EnqueueImageGenRequestSuccessResponse {
pub success: bool,
pub inference_job_token: InferenceJobToken,
}
#[derive(Debug, ToSchema)]
pub enum EnqueueImageGenRequestError {
BadInput(String),
NotAuthorized,
ServerError,
RateLimited,
}
impl ResponseError for EnqueueImageGenRequestError {
fn status_code(&self) -> StatusCode {
match *self {
EnqueueImageGenRequestError::BadInput(_) => StatusCode::BAD_REQUEST,
EnqueueImageGenRequestError::NotAuthorized => StatusCode::UNAUTHORIZED,
EnqueueImageGenRequestError::ServerError => StatusCode::INTERNAL_SERVER_ERROR,
EnqueueImageGenRequestError::RateLimited => StatusCode::TOO_MANY_REQUESTS,
}
}
fn error_response(&self) -> HttpResponse {
let error_reason = match self {
EnqueueImageGenRequestError::BadInput(reason) => reason.to_string(),
EnqueueImageGenRequestError::NotAuthorized => "unauthorized".to_string(),
EnqueueImageGenRequestError::ServerError => "server error".to_string(),
EnqueueImageGenRequestError::RateLimited => "rate limited".to_string(),
};
to_simple_json_error(&error_reason, self.status_code())
}
}
impl Display for EnqueueImageGenRequestError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{:?}", self)
}
}
/// Prompt image generation using image studio
#[utoipa::path(
post,
tag = "Image Studio",
path = "/v1/image_studio/prompt",
responses(
(
status = 200,
description = "Enqueue TTS generically",
body = EnqueueImageGenRequestSuccessResponse,
),
(status = 400, description = "Bad input", body = EnqueueImageGenRequestError),
(status = 401, description = "Not authorized", body = EnqueueImageGenRequestError),
(status = 429, description = "Rate limited", body = EnqueueImageGenRequestError),
(status = 500, description = "Server error", body = EnqueueImageGenRequestError)
),
params(("request" = EnqueueStudioImageGenRequest, description = "Payload for Image Generation Request"))
)]
pub async fn enqueue_studio_image_generation_handler(
http_request: HttpRequest,
request: Json<EnqueueStudioImageGenRequest>,
server_state: Data<Arc<ServerState>>
) -> Result<HttpResponse, EnqueueImageGenRequestError> {
validate_request(&request)?;
let mut mysql_connection = server_state.mysql_pool.acquire().await.map_err(|err| {
warn!("MySql pool error: {:?}", err);
EnqueueImageGenRequestError::ServerError
})?;
// ==================== USER SESSION ==================== //
let maybe_user_session = server_state.session_checker
.maybe_get_user_session_extended_from_connection(&http_request, &mut mysql_connection).await
.map_err(|e| {
warn!("Session checker error: {:?}", e);
EnqueueImageGenRequestError::ServerError
})?;
let mut maybe_user_token: Option<UserToken> = None;
if let Some(user_session) = maybe_user_session.as_ref() {
maybe_user_token = Some(UserToken::new_from_str(&user_session.user_token));
}
// ==================== PLANS ==================== //
// Plan should handle "first anonymous use" and "investor" cases.
let plan = get_correct_plan_for_session(
server_state.server_environment_old,
maybe_user_session.as_ref()
);
// Separate priority for animation.
let priority_level = plan.web_vc_base_priority_level();
// ==================== DEBUG MODE + ROUTING TAG ==================== //
let is_debug_request = get_request_header_optional(&http_request, DEBUG_HEADER_NAME).is_some();
let maybe_routing_tag = get_request_header_optional(&http_request, ROUTING_TAG_HEADER_NAME).map(
|routing_tag| routing_tag.trim().to_string()
);
// ==================== BANNED USERS ==================== //
if let Some(ref user) = maybe_user_session {
if user.role.is_banned {
return Err(EnqueueImageGenRequestError::NotAuthorized);
}
}
// DETECT premium user and queue
// ==================== RATE LIMIT ==================== //
let rate_limiter = match maybe_user_session {
None => &server_state.redis_rate_limiters.logged_out,
Some(ref _user) => &server_state.redis_rate_limiters.logged_in,
};
if let Err(_err) = rate_limiter.rate_limit_request(&http_request) {
return Err(EnqueueImageGenRequestError::RateLimited);
}
// Get up IP address
let ip_address = get_request_ip(&http_request);
// ==================== HANDLE IDEMPOTENCY ==================== //
if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) {
return Err(EnqueueImageGenRequestError::BadInput(reason));
}
insert_idempotency_token(&request.uuid_idempotency_token, &mut *mysql_connection)
.await
.map_err(|err| {
error!("Error inserting idempotency token: {:?}", err);
EnqueueImageGenRequestError::BadInput("invalid idempotency token".to_string())
})?;
// ==================== DOWNLOAD FILES AND UPLOAD TO SORA ==================== //
// TODO: Maybe this moves to a job.
let work_temp_dir = server_state.temp_dir_creator.new_tempdir("image_studio").unwrap();
let public_bucket_client = server_state.public_bucket_client.clone();
let scene_media_token = request.snapshot_media_token.clone();
let scene_media_path = query_and_download_media_file(
&scene_media_token,
&work_temp_dir,
&public_bucket_client,
&server_state.mysql_pool
).await.map_err(|err| {
error!("Failed to download scene media file: {:?}", err);
EnqueueImageGenRequestError::ServerError
})?;
let mut files_to_upload = Vec::with_capacity(1 + request.maybe_additional_images
.as_ref().map(|v| v.len()).unwrap_or(0));
files_to_upload.push(scene_media_path.clone());
let additional_images = request.maybe_additional_images.clone().unwrap_or_else(|| vec![]);
for media_file_token in &additional_images {
let media_file_path = query_and_download_media_file(
&media_file_token,
&work_temp_dir,
&public_bucket_client,
&server_state.mysql_pool
).await.map_err(|err| {
error!("Failed to download additional media file: {:?}", err);
EnqueueImageGenRequestError::ServerError
})?;
files_to_upload.push(media_file_path.clone());
}
// ==================== HANDLE SORA CREDENTIALS ==================== //
let mut redis = server_state.redis_pool
.get()
.map_err(|e| {
error!("redis error: {:?}", e);
EnqueueImageGenRequestError::ServerError
})?;
let sora_credentials = get_sora_credentials_from_request(&http_request, &mut redis)
.map_err(|e| {
error!("sora credential error: {:?}", e);
EnqueueImageGenRequestError::ServerError
})?;
// ==================== HANDLE SORA UPLOAD ==================== //
let mut sora_media_tokens = Vec::with_capacity(files_to_upload.len());
for file_path in files_to_upload {
let scene_media_upload_request = SoraMediaUploadRequest {
file_path: scene_media_path.to_string_lossy().to_string(),
credentials: &sora_credentials,
};
let sora_upload_response = sora_media_upload_from_file(file_path, &sora_credentials)
.await
.map_err(|err| {
error!("Failed to upload scene media to Sora: {:?}", err);
EnqueueImageGenRequestError::ServerError
})?;
debug!("Uploaded media to Sora : {:?}", sora_upload_response);
sora_media_tokens.push(sora_upload_response.id);
}
// ==================== HANDLE SORA PROMPT ==================== //
let number_of_samples = match request.maybe_number_of_samples {
Some(val) => if val > MAXIMUM_IMAGE_COUNT {
MAXIMUM_IMAGE_COUNT
} else if val < MINIMUM_IMAGE_COUNT {
MINIMUM_IMAGE_COUNT
} else {
val
},
None => 1,
};
let prompt = request.prompt.clone();
let response = sora_image_gen_remix(SoraImageGenRemixRequest {
prompt: prompt.clone(),
num_images: NumImages::One,
image_size: ImageSize::Square,
sora_media_tokens: sora_media_tokens.clone(),
credentials: &sora_credentials,
}).await
.map_err(|err| {
error!("Failed to call Sora image generation: {:?}", err);
EnqueueImageGenRequestError::ServerError
})?;
debug!("Sora image gen response: {:?}", response);
// Store the actual inference args that will go to the database
let inference_args = SoraImageGenArgs {
prompt: Some(prompt),
scene_snapshot_media_token: Some(request.snapshot_media_token.clone()),
maybe_additional_media_file_tokens: Some(additional_images),
maybe_number_of_samples: Some(number_of_samples),
maybe_sora_media_upload_tokens: Some(sora_media_tokens),
};
// create the inference args here
let maybe_avt_token = server_state.avt_cookie_manager.get_avt_token_from_request(&http_request);
// create the job record here!
let query_result = insert_generic_inference_job(InsertGenericInferenceArgs {
uuid_idempotency_token: &request.uuid_idempotency_token,
job_type: InferenceJobType::ImageGenApi,
maybe_product_category: None, // This is not a product anymore
inference_category: InferenceCategory::ImageGeneration,
maybe_model_type: Some(InferenceModelType::ImageGenApi), // NB: Model is static during inference
maybe_model_token: None, // NB: Model is static during inference
maybe_input_source_token: None,
maybe_input_source_token_type: None,
maybe_download_url: None,
maybe_cover_image_media_file_token: None,
maybe_raw_inference_text: None,
maybe_max_duration_seconds: None,
maybe_inference_args: Some(GenericInferenceArgs {
inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration),
args: Some(PolymorphicInferenceArgs::Sg(inference_args)),
}),
maybe_creator_user_token: maybe_user_token.as_ref(),
maybe_avt_token: maybe_avt_token.as_ref(),
creator_ip_address: &ip_address,
creator_set_visibility: Visibility::Public,
priority_level,
requires_keepalive: false, //reverse ... TODO fix this. we set it base on account is premium or not ...
is_debug_request,
maybe_routing_tag: maybe_routing_tag.as_deref(),
mysql_pool: &server_state.mysql_pool,
}).await;
let job_token = match query_result {
Ok((job_token, _id)) => job_token,
Err(err) => {
warn!("New generic inference job creation DB error: {:?}", err);
if err.had_duplicate_idempotency_token() {
return Err(EnqueueImageGenRequestError::BadInput("Duplicate idempotency token".to_string()));
}
return Err(EnqueueImageGenRequestError::ServerError);
}
};
let response: EnqueueImageGenRequestSuccessResponse = EnqueueImageGenRequestSuccessResponse {
success: true,
inference_job_token: job_token,
};
let body = serde_json::to_string(&response)
.map_err(|_e| EnqueueImageGenRequestError::ServerError)?;
// Error handling 101 rust result type returned like so.
Ok(HttpResponse::Ok().content_type("application/json").body(body))
}
fn validate_request(
request: &Json<EnqueueStudioImageGenRequest>,
) -> Result<(), EnqueueImageGenRequestError> {
Ok(())
}
#[derive(Debug)]
pub enum DownloadMediaFileError {
IoError(std::io::Error),
MediaFileNotFound,
Other(anyhow::Error),
}
impl From<anyhow::Error> for DownloadMediaFileError {
fn from(err: anyhow::Error) -> Self {
DownloadMediaFileError::Other(err)
}
}
async fn query_and_download_media_file(
media_file_token: &MediaFileToken,
download_dir: &TempDir,
bucket_client: &BucketClient,
mysql_pool: &MySqlPool,
) -> Result<PathBuf, DownloadMediaFileError> {
let media_file_result = get_media_file(
&media_file_token,
false,
mysql_pool
).await;
let media_file= match media_file_result {
Ok(Some(result)) => result,
Ok(None) => {
warn!("could not find media file in database: {:?}", media_file_token);
return Err(DownloadMediaFileError::MediaFileNotFound)
}
Err(e) => {
error!("could not query media file: {:?}", e);
return Err(DownloadMediaFileError::Other(e))
}
};
let media_file_bucket_path = MediaFileBucketPath::from_object_hash(
&media_file.public_bucket_directory_hash,
media_file.maybe_public_bucket_prefix.as_deref(),
media_file.maybe_public_bucket_extension.as_deref());
let extension = media_file.maybe_public_bucket_extension.clone().unwrap_or_else(|| ".jpg".to_string());
let download_filename = format!("download_file_{}.{}", media_file.token.as_str(), extension);
let download_file_path = download_dir.path().join(download_filename);
info!("Downloading from bucket...");
bucket_client.download_file_to_disk(
&media_file_bucket_path.to_full_object_pathbuf(),
&download_file_path,
).await?;
Ok(download_file_path)
}
@@ -0,0 +1 @@
pub mod enqueue_studio_image_generation_handler;
@@ -0,0 +1 @@
pub mod upload_snapshot_media_file_handler;
@@ -0,0 +1,279 @@
use std::collections::HashSet;
use std::io::{BufReader, Bytes, Cursor, Read};
use std::path::PathBuf;
use std::sync::Arc;
use actix_multipart::form::tempfile::TempFile;
use actix_multipart::form::text::Text;
use actix_multipart::form::MultipartForm;
use actix_multipart::Multipart;
use actix_web::web::Json;
use actix_web::{web, HttpRequest, HttpResponse};
use log::{error, info, warn};
use once_cell::sync::Lazy;
use utoipa::ToSchema;
use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
use enums::by_table::media_files::media_file_class::MediaFileClass;
use enums::by_table::media_files::media_file_engine_category::MediaFileEngineCategory;
use enums::by_table::media_files::media_file_type::MediaFileType;
use enums::common::visibility::Visibility;
use enums::no_table::style_transfer::style_transfer_name::StyleTransferName;
use hashing::sha256::sha256_hash_bytes::sha256_hash_bytes;
use http_server_common::request::get_request_ip::get_request_ip;
use media::decode_basic_audio_info::decode_basic_audio_bytes_info;
use mimetypes::mimetype_for_bytes::get_mimetype_for_bytes;
use mimetypes::mimetype_to_extension::mimetype_to_extension;
use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
use mysql_queries::queries::media_files::create::specialized_insert::insert_media_file_from_file_upload::{insert_media_file_from_file_upload, InsertMediaFileFromUploadArgs, UploadType};
use mysql_queries::queries::media_files::get::get_media_file::get_media_file;
use tokens::tokens::media_files::MediaFileToken;
use crate::http_server::endpoints::media_files::upload::upload_error::MediaFileUploadError;
use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format;
use crate::state::server_state::ServerState;
use crate::util::check_creator_tokens::{check_creator_tokens, CheckCreatorTokenArgs, CheckCreatorTokenResult};
/// Form-multipart request fields.
///
/// IF VIEWING DOCS, PLEASE SEE BOTTOM OF PAGE `UploadSnapshotMediaFileForm` (Under "Schema") FOR DETAILS ON FIELDS AND NULLABILITY.
#[derive(MultipartForm, ToSchema)]
#[multipart(duplicate_field = "deny")]
pub struct UploadSnapshotMediaFileForm {
/// UUID for request idempotency
#[multipart(limit = "2 KiB")]
#[schema(value_type = String, format = Binary)]
uuid_idempotency_token: Text<String>,
/// The uploaded "screenshot" / "snapshot" of the engine scene
#[multipart(limit = "512 MiB")]
#[schema(value_type = Vec<u8>, format = Binary)]
snapshot: TempFile,
/// Optional: If an engine scene was used to create this snapshot, provide it here to create a link
#[multipart(limit = "2 KiB")]
#[schema(value_type = Option<MediaFileToken>, format = Binary)]
scene_media_token: Option<Text<MediaFileToken>>,
/// Optional: Visibility of the scene
#[multipart(limit = "2 KiB")]
#[schema(value_type = Option<Visibility>, format = Binary)]
maybe_visibility: Option<Text<Visibility>>,
}
// Unlike the "upload" endpoints, which are pure inserts, these endpoints are *upserts*.
#[derive(Serialize, ToSchema)]
pub struct UploadSnapshotMediaFileSuccessResponse {
pub success: bool,
/// The media token of the "screenshot" / "snapshot" we took.
pub snapshot_media_token: MediaFileToken,
}
static ALLOWED_MIME_TYPES : Lazy<HashSet<&'static str>> = Lazy::new(|| {
HashSet::from([
"image/jpeg",
"image/png",
"image/gif",
"image/webp",
])
});
/// Upload snapshots of the scenes created in the 3D image studio.
#[utoipa::path(
post,
tag = "Image Studio",
path = "/v1/image_studio/scene_snapshot",
responses(
(status = 200, description = "Success Update", body = UploadSnapshotMediaFileSuccessResponse),
(status = 400, description = "Bad input", body = MediaFileUploadError),
(status = 401, description = "Not authorized", body = MediaFileUploadError),
(status = 429, description = "Too many requests", body = MediaFileUploadError),
(status = 500, description = "Server error", body = MediaFileUploadError),
),
params(
(
"request" = UploadSnapshotMediaFileForm,
description = "IF VIEWING DOCS, PLEASE SEE BOTTOM OF PAGE `UploadSnapshotMediaFileForm` (Under 'Schema') FOR DETAILS ON FIELDS AND NULLABILITY."
),
)
)]
pub async fn upload_snapshot_media_file_handler(
http_request: HttpRequest,
server_state: web::Data<Arc<ServerState>>,
MultipartForm(mut form): MultipartForm<UploadSnapshotMediaFileForm>,
) -> Result<Json<UploadSnapshotMediaFileSuccessResponse>, MediaFileUploadError> {
let mut mysql_connection = server_state.mysql_pool
.acquire()
.await
.map_err(|err| {
error!("MySql pool error: {:?}", err);
MediaFileUploadError::ServerError
})?;
// ==================== READ SESSION ==================== //
let maybe_user_session = server_state
.session_checker
.maybe_get_user_session_from_connection(&http_request, &mut mysql_connection)
.await
.map_err(|e| {
error!("Session checker error: {:?}", e);
MediaFileUploadError::ServerError
})?;
let maybe_user_token = maybe_user_session
.as_ref()
.map(|session| session.get_strongly_typed_user_token());
let maybe_avt_token = server_state
.avt_cookie_manager
.get_avt_token_from_request(&http_request);
// ==================== BANNED USERS ==================== //
if let Some(ref user) = maybe_user_session {
if user.is_banned {
return Err(MediaFileUploadError::NotAuthorizedVerbose("user is banned".to_string()));
}
}
// ==================== HANDLE IDEMPOTENCY ==================== //
// TODO(bt, 2024-02-26): This should be a transaction.
let uuid_idempotency_token = form.uuid_idempotency_token.as_ref();
if let Err(reason) = validate_idempotency_token_format(uuid_idempotency_token) {
return Err(MediaFileUploadError::BadInput(reason));
}
insert_idempotency_token(uuid_idempotency_token, &mut *mysql_connection)
.await
.map_err(|err| {
error!("Error inserting idempotency token: {:?}", err);
MediaFileUploadError::BadInput("invalid idempotency token".to_string())
})?;
// ==================== UPLOAD METADATA ==================== //
let creator_set_visibility = form.maybe_visibility
.map(|visibility| visibility.0)
.or_else(|| {
maybe_user_session
.as_ref()
.map(|user_session| user_session.preferred_tts_result_visibility)
})
.unwrap_or(Visibility::default());
// ==================== USER DATA ==================== //
let ip_address = get_request_ip(&http_request);
// ==================== FILE VALIDATION ==================== //
let mut file_bytes = Vec::new();
form.snapshot.file.read_to_end(&mut file_bytes)
.map_err(|e| {
error!("Problem reading file: {:?}", e);
MediaFileUploadError::ServerError
})?;
let mimetype = get_mimetype_for_bytes(file_bytes.as_ref())
.map(|mimetype| mimetype.to_string())
.ok_or_else(|| {
warn!("Could not determine mimetype for file");
MediaFileUploadError::BadInput("Could not determine mimetype for file".to_string())
})?;
if !ALLOWED_MIME_TYPES.contains(mimetype.as_str()) {
// NB: Don't let our error message inject malicious strings
let filtered_mimetype = mimetype
.chars()
.filter(|c| c.is_ascii())
.filter(|c| c.is_alphanumeric() || *c == '/')
.collect::<String>();
return Err(MediaFileUploadError::BadInput(format!("unpermitted mime type: {}", &filtered_mimetype)));
}
// ==================== OTHER FILE METADATA ==================== //
let maybe_filename = form.snapshot.file_name.as_deref()
.as_deref()
.map(|filename| PathBuf::from(filename));
let extension = mimetype_to_extension(&mimetype)
.or_else(|| {
maybe_filename
.as_ref()
.and_then(|filename| filename.extension()) // TODO needs dot prefix
.and_then(|ext| ext.to_str())
})
.ok_or_else(|| {
warn!("Could not determine file extension for mimetype: {}", &mimetype);
MediaFileUploadError::ServerError
})?;
let extension = format!(".{extension}"); // NB: needs dot prefix
let file_size_bytes = file_bytes.len();
let hash = sha256_hash_bytes(&file_bytes)
.map_err(|io_error| {
error!("Problem hashing bytes: {:?}", io_error);
MediaFileUploadError::ServerError
})?;
// ==================== UPLOAD AND SAVE ==================== //
const PREFIX : Option<&str> = Some("snapshot_");
let public_upload_path = MediaFileBucketPath::generate_new(PREFIX, Some(&extension));
info!("Uploading media to bucket path: {}", public_upload_path.get_full_object_path_str());
server_state.public_bucket_client.upload_file_with_content_type(
public_upload_path.get_full_object_path_str(),
file_bytes.as_ref(),
&mimetype)
.await
.map_err(|e| {
warn!("Upload media bytes to bucket error: {:?}", e);
MediaFileUploadError::ServerError
})?;
let (token, record_id) = insert_media_file_from_file_upload(InsertMediaFileFromUploadArgs {
maybe_media_class: Some(MediaFileClass::Image),
media_file_type: MediaFileType::Image,
maybe_creator_user_token: maybe_user_token.as_ref(),
maybe_creator_anonymous_visitor_token: maybe_avt_token.as_ref(),
creator_ip_address: &ip_address,
creator_set_visibility,
upload_type: UploadType::Filesystem, // TODO(bt,2024-05-02): This should be a parameter and a well-known enum.
maybe_engine_category: None,
maybe_animation_type: None,
maybe_mime_type: Some(&mimetype),
file_size_bytes: file_size_bytes as u64,
maybe_duration_millis: None,
sha256_checksum: &hash,
maybe_title: Some("Snapshot"),
maybe_scene_source_media_file_token: None,
is_intermediate_system_file: true,
public_bucket_directory_hash: public_upload_path.get_object_hash(),
maybe_public_bucket_prefix: PREFIX,
maybe_public_bucket_extension: Some(&extension),
pool: &server_state.mysql_pool,
})
.await
.map_err(|err| {
warn!("New file creation DB error: {:?}", err);
MediaFileUploadError::ServerError
})?;
info!("new media file id: {} token: {:?}", record_id, &token);
Ok(Json(UploadSnapshotMediaFileSuccessResponse {
success: true,
snapshot_media_token: token,
}))
}
@@ -5,6 +5,7 @@ pub mod app_state;
pub mod beta_keys;
pub mod categories;
pub mod comments;
pub mod control_plane;
pub mod conversion;
pub mod download_job;
pub mod engine;
@@ -12,6 +13,7 @@ pub mod events;
pub mod featured_items;
pub mod flags;
pub mod image_gen;
pub mod image_studio;
pub mod inference_job;
pub mod investor_demo;
pub mod leaderboard;
@@ -30,9 +32,9 @@ pub mod subscriptions;
pub mod tags;
pub mod trending;
pub mod tts;
pub mod users;
pub mod user_bookmarks;
pub mod user_ratings;
pub mod users;
pub mod vocoders;
pub mod voice_clone_requests;
pub mod voice_conversion;
@@ -0,0 +1,24 @@
use actix_http::body::MessageBody;
use actix_service::ServiceFactory;
use actix_web::dev::{ServiceRequest, ServiceResponse};
use actix_web::{web, App, Error, HttpResponse};
use crate::http_server::endpoints::control_plane::set_sora_secret_handler::set_sora_secret_handler;
pub fn add_control_plane_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/control_plane")
.service(web::resource("/sora_secret")
.route(web::post().to(set_sora_secret_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
)
}
@@ -0,0 +1,29 @@
use crate::http_server::endpoints::image_studio::prompt::enqueue_studio_image_generation_handler::enqueue_studio_image_generation_handler;
use crate::http_server::endpoints::image_studio::upload::upload_snapshot_media_file_handler::upload_snapshot_media_file_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_image_studio_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/image_studio")
.service(web::resource("/scene_snapshot")
.route(web::post().to(upload_snapshot_media_file_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(web::resource("/prompt")
.route(web::post().to(enqueue_studio_image_generation_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
)
}
@@ -123,6 +123,8 @@ use crate::http_server::endpoints::w2l::get_w2l_template_use_count::get_w2l_temp
use crate::http_server::endpoints::w2l::get_w2l_upload_template_job_status::get_w2l_upload_template_job_status_handler;
use crate::http_server::endpoints::w2l::list_w2l_templates::list_w2l_templates_handler;
use crate::http_server::endpoints::w2l::set_w2l_template_mod_approval::set_w2l_template_mod_approval_handler;
use crate::http_server::routes::add_control_plane_routes::add_control_plane_routes;
use crate::http_server::routes::add_image_studio_routes::add_image_studio_routes;
use crate::http_server::routes::add_studio_gen2_routes::add_studio_gen2_routes;
use crate::http_server::routes::beta_key_routes::add_beta_key_routes;
use crate::http_server::routes::job_routes::add_job_routes;
@@ -171,8 +173,10 @@ pub fn add_routes<T, B> (app: App<T>, server_environment: ServerEnvironment) ->
app = add_model_download_routes(app);
app = add_workflow_routes(app);
app = add_studio_gen2_routes(app);
app = add_image_studio_routes(app);
app = add_job_routes(app);
app = add_engine_routes(app); /* /v1/engine/... */
app = add_control_plane_routes(app); /* /v1/control_plane/... */
// app = add_image_gen_routes(app);
@@ -8,4 +8,6 @@ pub mod moderation_routes;
pub mod user_routes;
pub mod weights_routes;
pub mod workflow_routes;
mod tag_routes;
mod tag_routes;
mod add_image_studio_routes;
mod add_control_plane_routes;
@@ -11,6 +11,20 @@ pub struct ScopedTempDirCreator {
}
impl ScopedTempDirCreator {
pub fn auto_setup() -> Self {
// Production configuration
let maybe_tmp = easyenv::get_env_pathbuf_optional("TEMP_DIR");
// MacOS
// TODO: Only run this on macOS
let maybe_mac_tmp = easyenv::get_env_pathbuf_optional("TMPDIR");
let directory = maybe_tmp
.or(maybe_mac_tmp)
.unwrap_or_else(|| PathBuf::from("/tmp"));
Self::for_directory(directory)
}
pub fn for_directory<P: AsRef<Path>>(base_dir: P) -> Self {
Self {
@@ -464,7 +464,7 @@ async fn main() -> AnyhowResult<()> {
user_tokens: user_token_troll_bans,
ip_addresses: ip_address_troll_bans,
},
temp_dir_creator: ScopedTempDirCreator::for_directory(easyenv::get_env_pathbuf_or_default("TEMP_DIR", "/tmp")),
temp_dir_creator: ScopedTempDirCreator::auto_setup(),
google_sign_in_cert: GoogleSignInCert::new(),
};
@@ -11,5 +11,6 @@ pub mod enroll_in_studio;
pub mod generate_random_username;
pub mod placeholder_images;
pub mod read_toml_file_to_struct;
pub mod redis_sora_secrets;
pub mod title_to_url_slug;
pub mod troll_user_bans;
@@ -0,0 +1,67 @@
use std::collections::HashMap;
use anyhow::anyhow;
use r2d2_redis::r2d2::PooledConnection;
use r2d2_redis::redis::Commands;
use r2d2_redis::RedisConnectionManager;
use errors::AnyhowResult;
use openai_sora_client::credentials::SoraCredentials;
/// Redis key
const SORA_SECRET_REDIS_KEY : &str = "sora_secret";
// Fields within the HKEY
const BEARER_SUBKEY : &str = "bearer";
const COOKIE_SUBKEY : &str = "cookie";
const SENTINEL_SUBKEY : &str = "sentinel";
#[derive(Clone,Copy,Debug)]
pub enum RedisSoraCredentialSubkey {
Bearer,
Cookie,
Sentinel,
}
impl RedisSoraCredentialSubkey {
pub fn to_str(&self) -> &'static str {
match self {
RedisSoraCredentialSubkey::Bearer => BEARER_SUBKEY,
RedisSoraCredentialSubkey::Cookie => COOKIE_SUBKEY,
RedisSoraCredentialSubkey::Sentinel => SENTINEL_SUBKEY,
}
}
}
pub fn get_sora_credentials(
redis: &mut PooledConnection<RedisConnectionManager>
) -> AnyhowResult<SoraCredentials> {
let values : HashMap<String, String> = redis.hgetall(SORA_SECRET_REDIS_KEY)
.map_err(|e| anyhow!("Failed to get Sora credentials from Redis: {}", e))?;
let bearer = values.get(BEARER_SUBKEY);
let cookie = values.get(COOKIE_SUBKEY);
let sentinel = values.get(SENTINEL_SUBKEY);
match (bearer, cookie, sentinel) {
(Some(b), Some(c), Some(s)) => {
Ok(SoraCredentials {
bearer_token: b.to_string(),
cookie: c.to_string(),
sentinel: Some(s.to_string()),
})
}
_ => Err(anyhow!("redis sora credential values not present")),
}
}
pub fn set_sora_credential_field(
redis: &mut PooledConnection<RedisConnectionManager>,
subkey: RedisSoraCredentialSubkey,
value: &str,
) -> AnyhowResult<()> {
redis.hset(SORA_SECRET_REDIS_KEY, subkey.to_str(), value)
.map_err(|e| anyhow!("Failed to set Sora credentials in Redis: {}", e))?;
Ok(())
}
+248 -219
View File
@@ -15,252 +15,281 @@ use tokenizers::Tokenizer;
#[derive(Parser)]
#[command(author, version, about, long_about = None)]
struct Args {
/// The prompt to be used for image generation.
#[arg(long, default_value = "A rusty robot walking on a beach")]
prompt: String,
/// The prompt to be used for image generation.
#[arg(long, default_value = "A rusty robot walking on a beach")]
prompt: String,
/// Run on CPU rather than on GPU.
#[arg(long)]
cpu: bool,
/// Run on CPU rather than on GPU.
#[arg(long)]
cpu: bool,
/// Use the quantized model.
#[arg(long)]
quantized: bool,
/// Use the quantized model.
#[arg(long)]
quantized: bool,
/// Enable tracing (generates a trace-timestamp.json file).
#[arg(long)]
tracing: bool,
/// Enable tracing (generates a trace-timestamp.json file).
#[arg(long)]
tracing: bool,
/// The height in pixels of the generated image.
#[arg(long)]
height: Option<usize>,
/// The height in pixels of the generated image.
#[arg(long)]
height: Option<usize>,
/// The width in pixels of the generated image.
#[arg(long)]
width: Option<usize>,
/// The width in pixels of the generated image.
#[arg(long)]
width: Option<usize>,
#[arg(long)]
decode_only: Option<String>,
#[arg(long)]
decode_only: Option<String>,
#[arg(long, value_enum, default_value = "schnell")]
model: Model,
#[arg(long, value_enum, default_value = "schnell")]
model: Model,
/// Use the slower kernels.
#[arg(long)]
use_dmmv: bool,
/// Use the slower kernels.
#[arg(long)]
use_dmmv: bool,
/// The seed to use when generating random samples.
#[arg(long)]
seed: Option<u64>,
/// The seed to use when generating random samples.
#[arg(long)]
seed: Option<u64>,
#[arg(long)]
cache_path: Option<String>,
/// Number of blocks to keep active on GPU (default: 1).
/// Only relevant when NOT using --use_full_gpu.
/// Higher values use more memory but may improve performance.
#[arg(long, default_value = "1")]
active_blocks: usize,
/// Use CUDA stream for faster device transfers (default: true).
/// Only relevant when running on GPU.
#[arg(long, default_value = "true")]
use_cuda_stream: bool,
/// Enable prefetching of the next batch of blocks while processing the current batch.
/// Only relevant when NOT using --use_full_gpu.
/// This can improve performance by overlapping memory transfers and computation.
#[arg(long, default_value = "false")]
prefetch_next_batch: bool,
/// Load entire model on GPU instead of using memory-efficient active blocks approach.
/// This uses more GPU memory but avoids CPU-GPU transfers completely.
/// When this flag is enabled, the --active_blocks and --prefetch_next_batch flags have no effect.
#[arg(long, default_value = "false")]
use_full_gpu: bool,
}
#[derive(Debug, Clone, Copy, clap::ValueEnum, PartialEq, Eq)]
enum Model {
Schnell,
Dev,
Schnell,
Dev,
}
fn run(args: Args) -> Result<()> {
use tracing_chrome::ChromeLayerBuilder;
use tracing_subscriber::prelude::*;
use tracing_chrome::ChromeLayerBuilder;
use tracing_subscriber::prelude::*;
let Args {
prompt,
cpu,
height,
width,
tracing,
decode_only,
model,
quantized,
..
} = args;
let width = width.unwrap_or(1360);
let height = height.unwrap_or(768);
let Args { prompt, cpu, height, width, tracing, decode_only, model, quantized, active_blocks, use_cuda_stream, prefetch_next_batch, use_full_gpu, .. } = args;
let width = width.unwrap_or(1360);
let height = height.unwrap_or(768);
let _guard = if tracing {
let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
tracing_subscriber::registry().with(chrome_layer).init();
Some(guard)
} else {
None
let _guard = if tracing {
let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
tracing_subscriber::registry().with(chrome_layer).init();
Some(guard)
} else {
None
};
let api = match args.cache_path.as_ref() {
Some(path) => hf_hub::api::sync::ApiBuilder::from_cache(hf_hub::Cache::new(path.to_string().into())).build().map_err(anyhow::Error::msg)?,
None => {
let cache = hf_hub::Cache::from_env();
hf_hub::api::sync::ApiBuilder::from_cache(cache).build().map_err(anyhow::Error::msg)?
},
};
let bf_repo = {
let name = match model {
Model::Dev => "black-forest-labs/FLUX.1-dev",
Model::Schnell => "black-forest-labs/FLUX.1-schnell",
};
api.repo(hf_hub::Repo::model(name.to_string()))
};
let api = hf_hub::api::sync::Api::new()?;
let bf_repo = {
let name = match model {
Model::Dev => "black-forest-labs/FLUX.1-dev",
Model::Schnell => "black-forest-labs/FLUX.1-schnell",
};
api.repo(hf_hub::Repo::model(name.to_string()))
};
let device = candle_examples::device(cpu)?;
if let Some(seed) = args.seed {
device.set_seed(seed)?;
// Use CUDA with stream for faster device transfers when not in CPU mode
let device = if cpu {
candle::Device::Cpu
} else {
#[cfg(feature = "cuda")]
{
if use_cuda_stream {
println!("Using CUDA with dedicated stream for faster device transfers");
// Use device ordinal 0 (first GPU)
candle::Device::new_cuda_with_stream(0)?
} else {
println!("Using standard CUDA device (no dedicated stream)");
candle::Device::new_cuda(0)?
}
}
let dtype = device.bf16_default_to_f32();
let img = match decode_only {
None => {
let t5_emb = {
let repo = api.repo(hf_hub::Repo::with_revision(
"google/t5-v1_1-xxl".to_string(),
hf_hub::RepoType::Model,
"refs/pr/2".to_string(),
));
let model_file = repo.get("model.safetensors")?;
let vb =
unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
let config_filename = repo.get("config.json")?;
let config = std::fs::read_to_string(config_filename)?;
let config: t5::Config = serde_json::from_str(&config)?;
let mut model = t5::T5EncoderModel::load(vb, &config)?;
let tokenizer_filename = api
.model("lmz/mt5-tokenizers".to_string())
.get("t5-v1_1-xxl.tokenizer.json")?;
let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
let mut tokens = tokenizer
.encode(prompt.as_str(), true)
.map_err(E::msg)?
.get_ids()
.to_vec();
tokens.resize(256, 0);
let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
println!("{input_token_ids}");
model.forward(&input_token_ids)?
};
println!("T5\n{t5_emb}");
let clip_emb = {
let repo = api.repo(hf_hub::Repo::model(
"openai/clip-vit-large-patch14".to_string(),
));
let model_file = repo.get("model.safetensors")?;
let vb =
unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
// https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json
let config = clip::text_model::ClipTextConfig {
vocab_size: 49408,
projection_dim: 768,
activation: clip::text_model::Activation::QuickGelu,
intermediate_size: 3072,
embed_dim: 768,
max_position_embeddings: 77,
pad_with: None,
num_hidden_layers: 12,
num_attention_heads: 12,
};
let model =
clip::text_model::ClipTextTransformer::new(vb.pp("text_model"), &config)?;
let tokenizer_filename = repo.get("tokenizer.json")?;
let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
let tokens = tokenizer
.encode(prompt.as_str(), true)
.map_err(E::msg)?
.get_ids()
.to_vec();
let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
println!("{input_token_ids}");
model.forward(&input_token_ids)?
};
println!("CLIP\n{clip_emb}");
let img = {
let cfg = match model {
Model::Dev => flux::model::Config::dev(),
Model::Schnell => flux::model::Config::schnell(),
};
let img = flux::sampling::get_noise(1, height, width, &device)?.to_dtype(dtype)?;
let state = if quantized {
flux::sampling::State::new(
&t5_emb.to_dtype(candle::DType::F32)?,
&clip_emb.to_dtype(candle::DType::F32)?,
&img.to_dtype(candle::DType::F32)?,
)?
} else {
flux::sampling::State::new(&t5_emb, &clip_emb, &img)?
};
let timesteps = match model {
Model::Dev => {
flux::sampling::get_schedule(50, Some((state.img.dim(1)?, 0.5, 1.15)))
}
Model::Schnell => flux::sampling::get_schedule(4, None),
};
println!("{state:?}");
println!("{timesteps:?}");
if quantized {
let model_file = match model {
Model::Schnell => api
.repo(hf_hub::Repo::model("lmz/candle-flux".to_string()))
.get("flux1-schnell.gguf")?,
Model::Dev => todo!(),
};
let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(
model_file, &device,
)?;
#[cfg(not(feature = "cuda"))]
{
println!("CUDA not available, using CPU");
candle::Device::Cpu
}
};
let model = flux::quantized_model::Flux::new(&cfg, vb)?;
flux::sampling::denoise(
&model,
&state.img,
&state.img_ids,
&state.txt,
&state.txt_ids,
&state.vec,
&timesteps,
4.,
)?
.to_dtype(dtype)?
} else {
let model_file = match model {
Model::Schnell => bf_repo.get("flux1-schnell.safetensors")?,
Model::Dev => bf_repo.get("flux1-dev.safetensors")?,
};
let vb = unsafe {
VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)?
};
let model = flux::model::Flux::new(&cfg, vb)?;
flux::sampling::denoise(
&model,
&state.img,
&state.img_ids,
&state.txt,
&state.txt_ids,
&state.vec,
&timesteps,
4.,
)?
}
};
flux::sampling::unpack(&img, height, width)?
}
Some(file) => {
let mut st = candle::safetensors::load(file, &device)?;
st.remove("img").unwrap().to_dtype(dtype)?
}
};
println!("latent img\n{img}");
let img = {
let model_file = bf_repo.get("ae.safetensors")?;
if let Some(seed) = args.seed {
device.set_seed(seed)?;
}
let start_time = std::time::Instant::now();
let dtype = device.bf16_default_to_f32();
let img = match decode_only {
None => {
let t5_emb = {
let repo = api.repo(hf_hub::Repo::with_revision("google/t5-v1_1-xxl".to_string(), hf_hub::RepoType::Model, "refs/pr/2".to_string()));
let model_file = repo.get("model.safetensors")?;
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
let config_filename = repo.get("config.json")?;
let config = std::fs::read_to_string(config_filename)?;
let config: t5::Config = serde_json::from_str(&config)?;
let mut model = t5::T5EncoderModel::load(vb, &config)?;
let tokenizer_filename = api.model("lmz/mt5-tokenizers".to_string()).get("t5-v1_1-xxl.tokenizer.json")?;
let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
let mut tokens = tokenizer.encode(prompt.as_str(), true).map_err(E::msg)?.get_ids().to_vec();
tokens.resize(256, 0);
let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
println!("{input_token_ids}");
model.forward(&input_token_ids)?
};
println!("T5\n{t5_emb}");
let clip_emb = {
let repo = api.repo(hf_hub::Repo::model("openai/clip-vit-large-patch14".to_string()));
let model_file = repo.get("model.safetensors")?;
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
// https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json
let config = clip::text_model::ClipTextConfig { vocab_size: 49408, projection_dim: 768, activation: clip::text_model::Activation::QuickGelu, intermediate_size: 3072, embed_dim: 768, max_position_embeddings: 77, pad_with: None, num_hidden_layers: 12, num_attention_heads: 12 };
let model = clip::text_model::ClipTextTransformer::new(vb.pp("text_model"), &config)?;
let tokenizer_filename = repo.get("tokenizer.json")?;
let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
let tokens = tokenizer.encode(prompt.as_str(), true).map_err(E::msg)?.get_ids().to_vec();
let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
println!("{input_token_ids}");
model.forward(&input_token_ids)?
};
println!("CLIP\n{clip_emb}");
let img = {
let cfg = match model {
Model::Dev => flux::autoencoder::Config::dev(),
Model::Schnell => flux::autoencoder::Config::schnell(),
Model::Dev => flux::model::Config::dev(),
Model::Schnell => flux::model::Config::schnell(),
};
let model = flux::autoencoder::AutoEncoder::new(&cfg, vb)?;
model.decode(&img)?
let img = flux::sampling::get_noise(1, height, width, &device)?.to_dtype(dtype)?;
let state = if quantized { flux::sampling::State::new(&t5_emb.to_dtype(candle::DType::F32)?, &clip_emb.to_dtype(candle::DType::F32)?, &img.to_dtype(candle::DType::F32)?)? } else { flux::sampling::State::new(&t5_emb, &clip_emb, &img)? };
let timesteps = match model {
Model::Dev => flux::sampling::get_schedule(28, Some((state.img.dim(1)?, 0.5, 1.15))),
Model::Schnell => flux::sampling::get_schedule(4, None),
};
println!("{state:?}");
println!("{timesteps:?}");
if quantized {
// Note about quantized model memory management
let model_file = match model {
Model::Schnell => api.repo(hf_hub::Repo::model("lmz/candle-flux".to_string())).get("flux1-schnell.gguf")?,
Model::Dev => todo!(),
};
let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(model_file, &device)?;
let mut model = flux::quantized_model::Flux::new(&cfg, vb)?;
// Quantized models are always fully loaded to the selected device
if !cpu {
println!("Quantized model fully loaded to GPU with optimized memory layout");
}
flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, &timesteps, 4.)?.to_dtype(dtype)?
} else {
let model_file = match model {
Model::Schnell => bf_repo.get("flux1-schnell.safetensors")?,
Model::Dev => bf_repo.get("flux1-dev.safetensors")?,
};
let mut model = if use_full_gpu {
// When use_full_gpu is true, load everything directly to GPU
println!("Loading entire model directly to GPU...");
let gpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
let mut model = flux::model::Flux::new(&cfg, gpu_vb)?;
model.set_use_device_management(false);
println!("Model loaded entirely on GPU");
model
} else {
// Use memory-efficient approach with blocks on CPU
// Create a CPU device for initially loading all model weights
let cpu_device = candle::Device::Cpu;
// Load the entire model to CPU first (uses less memory)
println!("Loading model to CPU to minimize GPU memory usage...");
let cpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &cpu_device)? };
// Create the model with CPU weights, then selectively move core components to GPU
let mut model = flux::model::Flux::new_with_gpu_core(&cfg, cpu_vb, &device)?;
println!("Model loaded (core components on GPU, blocks on CPU for reduced memory usage)");
// Set the number of active blocks to keep on GPU
model.set_max_active_blocks(active_blocks);
println!("Using {} active blocks on GPU", active_blocks);
// Enable prefetching if requested
if prefetch_next_batch {
model.set_prefetch_next_batch(true);
println!("Prefetching enabled - will overlap transfers with computation");
}
model
};
// Print final memory management approach being used
if use_full_gpu {
println!("Using full GPU memory - entire model on GPU");
} else {
println!("Using device management - blocks will be transferred between CPU and GPU as needed");
}
flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, &timesteps, 4.)?
}
};
flux::sampling::unpack(&img, height, width)?
},
Some(file) => {
let mut st = candle::safetensors::load(file, &device)?;
st.remove("img").unwrap().to_dtype(dtype)?
},
};
println!("latent img\n{img}");
let end_time = std::time::Instant::now();
println!("Time taken to decode: {:?}", end_time.duration_since(start_time));
let img = {
let model_file = bf_repo.get("ae.safetensors")?;
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
let cfg = match model {
Model::Dev => flux::autoencoder::Config::dev(),
Model::Schnell => flux::autoencoder::Config::schnell(),
};
println!("img\n{img}");
let img = ((img.clamp(-1f32, 1f32)? + 1.0)? * 127.5)?.to_dtype(candle::DType::U8)?;
let filename = match args.seed {
None => "out.jpg".to_string(),
Some(s) => format!("out-{s}.jpg"),
};
candle_examples::save_image(&img.i(0)?, filename)?;
Ok(())
let model = flux::autoencoder::AutoEncoder::new(&cfg, vb)?;
model.decode(&img)?
};
println!("img\n{img}");
let img = ((img.clamp(-1f32, 1f32)? + 1.0)? * 127.5)?.to_dtype(candle::DType::U8)?;
let filename = match args.seed {
None => "out.jpg".to_string(),
Some(s) => format!("out-{s}.jpg"),
};
candle_examples::save_image(&img.i(0)?, filename)?;
Ok(())
}
fn main() -> Result<()> {
let args = Args::parse();
#[cfg(feature = "cuda")]
candle::quantized::cuda::set_force_dmmv(args.use_dmmv);
run(args)
let args = Args::parse();
#[cfg(feature = "cuda")]
candle::quantized::cuda::set_force_dmmv(args.use_dmmv);
run(args)
}
+19 -1
View File
@@ -28,7 +28,8 @@
//! ```
//!
//! [`Layer Normalization`]: https://arxiv.org/abs/1607.06450
use candle::{DType, Module, Result, Tensor, D};
use candle::{DType, Module, Result, Tensor, D, Device};
use crate::offload::DeviceTransferable;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct LayerNormConfig {
@@ -201,3 +202,20 @@ pub fn rms_norm(size: usize, eps: f64, vb: crate::VarBuilder) -> Result<RmsNorm>
};
Ok(RmsNorm(layer_norm(size, config, vb)?))
}
impl DeviceTransferable for LayerNorm {
fn to_device(&self, device: &Device) -> Result<Self> {
Ok(Self {
weight: self.weight.to_device(device)?,
bias: self.bias.as_ref().map(|b| b.to_device(device)).transpose()?,
remove_mean: self.remove_mean,
eps: self.eps,
})
}
}
impl DeviceTransferable for RmsNorm {
fn to_device(&self, device: &Device) -> Result<Self> {
Ok(Self(self.0.to_device(device)?))
}
}
+2
View File
@@ -27,6 +27,7 @@ pub mod kv_cache;
pub mod layer_norm;
pub mod linear;
pub mod loss;
pub mod offload;
pub mod ops;
pub mod optim;
pub mod rnn;
@@ -58,3 +59,4 @@ pub use var_builder::VarBuilder;
pub use var_map::VarMap;
pub use candle::{Module, ModuleT};
pub use offload::DeviceTransferable;
+12 -1
View File
@@ -17,7 +17,9 @@
//! assert_eq!(ys.to_vec2::<f32>()?, &[[210.0, 430.0, 650.0]]);
//! # Ok(()) }
//! ```
use candle::{Result, Tensor};
use candle::{Result, Tensor, Device};
use crate::offload::DeviceTransferable;
#[derive(Clone, Debug)]
pub struct Linear {
@@ -88,3 +90,12 @@ pub fn linear_b(
linear_no_bias(in_dim, out_dim, vb)
}
}
impl DeviceTransferable for Linear {
fn to_device(&self, device: &Device) -> Result<Self> {
Ok(Self {
weight: self.weight.to_device(device)?,
bias: self.bias.as_ref().map(|b| b.to_device(device)).transpose()?,
})
}
}
+18
View File
@@ -0,0 +1,18 @@
use candle::{Device, Result};
pub trait DeviceTransferable {
/// Move the module to the specified device
///
/// # Arguments
///
/// * `device` - The target device to move the module to
///
/// # Returns
///
/// A new instance of the module with all tensors moved to the target device,
/// or an error if the module cannot be moved.
fn to_device(&self, _device: &Device) -> Result<Self> where Self: Sized {
// Default implementation that rejects moving - components can override this
candle::bail!("to_device not implemented for this type")
}
}
+2 -1
View File
@@ -9,7 +9,9 @@ candle = { path = "../candle-core", package = "candle-core" }
candle-flash-attn = { path = "../candle-flash-attn", optional = true }
candle-nn = { path = "../candle-nn", package = "candle-nn" }
fancy-regex = "0.13.0"
half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
intel-mkl-src = { version = "0.8.1", features = ["mkl-static-lp64-iomp"], optional = true }
log = "0.4.27"
num-traits = "0.2.15"
rand = "0.8.5"
rayon = "1.7.0"
@@ -17,7 +19,6 @@ serde = { version = "1.0.171", features = ["derive"] }
serde_json = "1.0.99"
serde_plain = "1.0.2"
tracing = "0.1.37"
half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
[features]
default = []
@@ -24,9 +24,39 @@
use candle::{Result, Tensor};
pub trait WithForward {
/// The forward pass for the Flux model.
///
/// # Arguments
///
/// * `img` - The input image tensor
/// * `img_ids` - The image position ids
/// * `txt` - The text tensor
/// * `txt_ids` - The text position ids
/// * `timesteps` - The timesteps tensor
/// * `y` - The conditioning vector
/// * `guidance` - An optional guidance vector
///
/// # Returns
///
/// The result of the forward pass
///
/// # Note on Memory Optimization
///
/// To optimize memory usage, we could implement a mechanism to move model blocks
/// to CPU when not in use. This would require:
///
/// 1. Adding `to_device` methods to all model blocks (DoubleStreamBlock, SingleStreamBlock, etc.)
/// 2. Implementing these methods by moving all tensors inside each block to the target device
/// 3. Modifying the forward method to:
/// - Move each block to GPU when needed
/// - Process the block
/// - Move the block back to CPU when done
///
/// This pattern could significantly reduce GPU memory usage for large models by
/// keeping only the currently active block in GPU memory.
#[allow(clippy::too_many_arguments)]
fn forward(
&self,
&mut self,
img: &Tensor,
img_ids: &Tensor,
txt: &Tensor,
File diff suppressed because it is too large Load Diff
@@ -419,7 +419,7 @@ impl Flux {
impl super::WithForward for Flux {
#[allow(clippy::too_many_arguments)]
fn forward(
&self,
&mut self,
img: &Tensor,
img_ids: &Tensor,
txt: &Tensor,
@@ -93,7 +93,7 @@ pub fn unpack(xs: &Tensor, height: usize, width: usize) -> Result<Tensor> {
#[allow(clippy::too_many_arguments)]
pub fn denoise<M: super::WithForward>(
model: &M,
model: &mut M,
img: &Tensor,
img_ids: &Tensor,
txt: &Tensor,
@@ -1,6 +1,9 @@
use super::schedulers::{betas_for_alpha_bar, BetaSchedule, PredictionType};
use candle::{DType, Device, IndexOp, Result, Tensor};
use log::debug;
const FINAL_ALPHA_CUMPROD: f64 = 0.9991499781608582;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LCMVarianceType {
FixedSmall,
@@ -87,7 +90,7 @@ impl LCMScheduler {
let timesteps = LCMScheduler::get_timesteps_for_steps(inference_steps, strength);
let step_ratio = config.train_timesteps / inference_steps;
println!("Using LCM timesteps: {:?}", timesteps);
debug!("Using LCM timesteps: {:?}", timesteps);
Ok(Self {
alphas_cumprod,
@@ -143,11 +146,11 @@ impl LCMScheduler {
let scaled_timestep_squared = scaled_timestep.powi(2);
let denominator = scaled_timestep_squared + sigma_squared;
println!(
debug!(
"Boundary condition for timestep {}, scaled timestep {}",
timestep, scaled_timestep
);
println!(
debug!(
" Intermediate values: sigma_data^2: {:.20}, scaled_timestep^2: {:.20}, denominator: {:.20}",
sigma_squared, scaled_timestep_squared, denominator
);
@@ -155,7 +158,7 @@ impl LCMScheduler {
let c_skip = sigma_squared / denominator;
let c_out = scaled_timestep / (scaled_timestep_squared + sigma_squared).sqrt();
println!(
debug!(
" Final computed values: c_skip: {:.20}, c_out: {:.20}",
c_skip, c_out
);
@@ -172,14 +175,14 @@ impl LCMScheduler {
) -> Result<Tensor> {
// Scale by exactly 1000.0 to match diffusers
let w = guidance_scale * 1000.0;
println!(
debug!(
"Creating guidance embedding with dim={}, scaled value={}",
embedding_dim, w
);
// Calculate the half dimension
let half_dim = embedding_dim / 2;
println!("Half dimension: {}", half_dim);
debug!("Half dimension: {}", half_dim);
let log_10000 = (10000.0f32).ln();
let freq_factor = log_10000 / ((half_dim - 1) as f32); // pre-compute division
@@ -191,7 +194,7 @@ impl LCMScheduler {
freqs.push(freq);
}
println!(
debug!(
"First few frequencies: {:?}",
&freqs.iter().take(5).collect::<Vec<_>>()
);
@@ -204,39 +207,39 @@ impl LCMScheduler {
emb_data.push(w_val * freq);
}
println!(
debug!(
"First few scaled frequencies: {:?}",
&emb_data.iter().take(5).collect::<Vec<_>>()
);
// Create the embedding tensor
let emb = Tensor::new(emb_data.as_slice(), device)?.to_dtype(dtype)?;
println!("Embedding tensor shape: {:?}", emb.shape());
debug!("Embedding tensor shape: {:?}", emb.shape());
// Reshape for correct sin/cos application
let emb = emb.reshape((1, half_dim))?;
println!("Reshaped embedding tensor: {:?}", emb.shape());
debug!("Reshaped embedding tensor: {:?}", emb.shape());
// Compute sin and cos components
let sin_emb = emb.sin()?;
println!("Sin embedding shape: {:?}", sin_emb.shape());
debug!("Sin embedding shape: {:?}", sin_emb.shape());
let cos_emb = emb.cos()?;
println!("Cos embedding shape: {:?}", cos_emb.shape());
debug!("Cos embedding shape: {:?}", cos_emb.shape());
// Concatenate sin and cos (along dimension 1)
let result = Tensor::cat(&[sin_emb, cos_emb], 1)?;
println!("After concatenation shape: {:?}", result.shape());
debug!("After concatenation shape: {:?}", result.shape());
// If odd dimension, pad with zero
let result = if embedding_dim % 2 == 1 {
println!("Adding zero padding for odd dimension");
debug!("Adding zero padding for odd dimension");
let zero = Tensor::zeros((1, 1), dtype, device)?;
Tensor::cat(&[result, zero], 1)?
} else {
result
};
println!("Final guidance embedding shape: {:?}", result.shape());
debug!("Final guidance embedding shape: {:?}", result.shape());
Ok(result)
}
@@ -288,12 +291,12 @@ impl LCMScheduler {
};
let timesteps = linspace.iter().map(|&f| f as usize).rev().collect();
println!("timesteps: {:?}", timesteps);
debug!("timesteps: {:?}", timesteps);
timesteps
}
};
println!("all_timesteps: {:?}", all_timesteps);
debug!("all_timesteps: {:?}", all_timesteps);
// Compute `t_start` based on `strength`
let num_train_timesteps = 1000;
@@ -303,7 +306,7 @@ impl LCMScheduler {
.position(|&t| t <= init_timestep)
.unwrap_or(0);
println!(
debug!(
"get_timesteps: t_start: {} init_timestep: {} strength: {}",
t_start, init_timestep, strength
);
@@ -327,7 +330,7 @@ impl LCMScheduler {
0
};
println!(
debug!(
"LCM step: timestep={}, prev_timestep={}, step_index={}",
timestep, prev_timestep, step_index
);
@@ -357,15 +360,15 @@ impl LCMScheduler {
let denoised = ((c_out * pred_original_sample )? + (c_skip * sample)?)?;
let pred_prev_sample = if step_index != self.timesteps.len() - 1 {
println!("Adding noise at intermediate step {}", timestep);
debug!("Adding noise at intermediate step {}", timestep);
let noise = Tensor::randn_like(&denoised, 0.0, 1.0)?;
((alpha_prod_t_prev.sqrt() * denoised)? + (beta_prod_t_prev.sqrt() * noise)?)?
} else {
println!("No noise for final timestep {}", timestep);
debug!("No noise for final timestep {}", timestep);
denoised
};
println!(
debug!(
"Denoised sample shape: {:?}, first few values: {}",
pred_prev_sample.shape(),
pred_prev_sample
@@ -388,7 +391,7 @@ impl LCMScheduler {
// self.timesteps[step_index + 1]
// };
// println!(
// debug!(
// "LCM step: timestep={}, prev_timestep={}",
// timestep, prev_timestep
// );
@@ -401,7 +404,7 @@ impl LCMScheduler {
// self.alphas_cumprod[prev_timestep]
// };
// println!(
// debug!(
// "LCM scheduler step: timestep={}, step_index={}, sigma={}",
// timestep,
// step_index,
@@ -433,13 +436,13 @@ impl LCMScheduler {
// pred_original_sample = pred_original_sample.clamp(-1f32, 1f32)?;
// }
// println!("boundary: Pred original sample first few values: {}",
// debug!("boundary: Pred original sample first few values: {}",
// pred_original_sample.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
// // Denoise model output using boundary conditions
// let denoised = ((pred_original_sample * c_out)? + (sample * c_skip)?)?;
// println!("boundary: Denoised first few values: {}",
// debug!("boundary: Denoised first few values: {}",
// denoised.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
// // LCM doesn't add noise during inference
@@ -452,11 +455,11 @@ impl LCMScheduler {
// // let pred_prev_sample = pred_prev_sample.clamp(-1f32, 1f32)?;
// println!(
// debug!(
// "LCM scheduler pred_original shape: {:?}",
// pred_prev_sample.shape()
// );
// println!(
// debug!(
// "LCM scheduler pred_original first few values: {}",
// pred_prev_sample
// .i(0)?
@@ -469,7 +472,7 @@ impl LCMScheduler {
// }
pub fn add_noise(&self, original: &Tensor, noise: Tensor, timestep: usize) -> Result<Tensor> {
println!("Adding noise at timestep {}", timestep);
debug!("Adding noise at timestep {}", timestep);
let timestep = if timestep >= self.alphas_cumprod.len() {
self.alphas_cumprod.len() - 1 // Ensure valid range
} else {
@@ -491,7 +494,7 @@ impl LCMScheduler {
// let alpha_cumprod = self.alphas_cumprod[timestep];
// // Print the alpha value to debug
// println!(
// debug!(
// "Adding noise with timestep={}, alpha={}, sqrt(alpha)={}",
// timestep,
// alpha_cumprod,
File diff suppressed because it is too large Load Diff
+26
View File
@@ -0,0 +1,26 @@
# Logs
logs
*.log
npm-debug.log*
yarn-debug.log*
yarn-error.log*
pnpm-debug.log*
lerna-debug.log*
node_modules
dist
dist-ssr
*.local
backup
.venv
# Editor directories and files
~
.vscode
.idea
.DS_Store
*.suo
*.ntvs*
*.njsproj
*.sln
*.sw?
**/*.onnx
+8
View File
@@ -0,0 +1,8 @@
<html>
<head>
<title>Not Found</title>
</head>
<body>
<h1> This Page Was Not Found </h1>
</body>
</html>
+39
View File
@@ -0,0 +1,39 @@
# Memeboard
## Installation & Dev
```sh
npm install
npm run dev
```
or
```sh
yarn # this installs the node packages
yarn dev # this runs the vite dev server
```
You may need to turn off browser security to get around CORS restrictions:
### Windows 11
from your powershell
```sh
# Chrome
Start "C:\Program Files\Google\Chrome\Application\chrome.exe" -ArgumentList '--user-data-dir="C://chrome-dev-disabled-security" --disable-web-security --disable-site-isolation-trials'
```
### on MAC
form your terminal
```sh
# Chrome
/Applications/Google\ Chrome.app/Contents/MacOS/Google\ Chrome --disable-web-security --user-data-dir=~/chrome
```
### Inspecting the Shared Workers
on chrome, type `chrome://inspect/#workers` into the url field
+3
View File
@@ -0,0 +1,3 @@
/*
Access-Control-Allow-Origin: *
+11
View File
@@ -0,0 +1,11 @@
# Board Netlify Redirects File
# Netlify docs:
# The redirects engine will process the first matching rule it
# finds, reading from top to bottom. Rules in the _redirects file
# are always processed first, followed by rules in the Netlify
# configuration file.
/landing /index.html 200!
/login /index.html 200!
/m_* /index.html 200!
+44
View File
@@ -0,0 +1,44 @@
import js from "@eslint/js";
import globals from "globals";
import reactHooks from "eslint-plugin-react-hooks";
import reactRefresh from "eslint-plugin-react-refresh";
import tseslint from "typescript-eslint";
import react from "eslint-plugin-react";
export default tseslint.config(
{ ignores: ["dist"] },
{
extends: [
js.configs.recommended,
...tseslint.configs.recommendedTypeChecked,
...tseslint.configs.stylisticTypeChecked,
],
files: ["**/*.{ts,tsx}"],
languageOptions: {
ecmaVersion: 2020,
globals: globals.browser,
},
plugins: {
"react-hooks": reactHooks,
"react-refresh": reactRefresh,
react,
},
rules: {
...reactHooks.configs.recommended.rules,
"react-refresh/only-export-components": [
"warn",
{ allowConstantExport: true },
],
...react.configs.recommended.rules,
...react.configs["jsx-runtime"].rules,
// TODO WIL: make no-reaching rule for imports
// "import/no-internal-modules": [
// "error",
// {
// allow: [],
// forbid: ["**/uiAccess/*"],
// },
// ],
},
},
);
+24
View File
@@ -0,0 +1,24 @@
<!doctype html>
<html lang="en">
<head>
<meta charset="UTF-8" />
<link rel="icon" type="image/x-ico" href="/favicon.ico" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<link rel="preconnect" href="https://fonts.googleapis.com" />
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin />
<link
href="https://fonts.googleapis.com/css2?family=Source+Sans+3:ital,wght@0,200..900;1,200..900&display=swap"
rel="stylesheet"
/>
<link
href="https://fonts.googleapis.com/css2?family=Syne:wght@600&display=swap"
rel="stylesheet"
/>
<title>Board | Storyteller.AI</title>
</head>
<body>
<div id="root"></div>
<script type="module" src="/src/app/index.tsx"></script>
</body>
</html>
+34
View File
@@ -0,0 +1,34 @@
[build]
command = "./script/netlify_build.sh"
publish = "dist/"
#[build.environment]
#REACT_APP_PUBLIC_POSTHOG_KEY = "phc_jBFgac0mVALAFk3negnSfYCcHgvkT00yBLQDmCDYNBb"
#REACT_APP_PUBLIC_POSTHOG_UI = "https://us.i.posthog.com"
#[dev]
#command = "npm run dev"
#framework = "vite"
# Set immutable caching for static files, because they have fingerprinted filenames
[[headers]]
for = "/build/*"
[headers.values]
"Cache-Control" = "public, max-age=31560000, immutable"
# [[redirects]]
# from = "/ingest/static/*"
# to = "https://us-assets.i.posthog.com/static/:splat"
# host = "us-assets.i.posthog.com"
# status = 200
# force = true
#
# [[redirects]]
# from = "/ingest/*"
# to = "https://us.i.posthog.com/:splat"
# host = "us.i.posthog.com"
# status = 200
# force = true
+9348
View File
File diff suppressed because it is too large Load Diff
+66
View File
@@ -0,0 +1,66 @@
{
"name": "storyteller-board",
"private": true,
"version": "0.0.0",
"type": "module",
"scripts": {
"dev": "echo 'realtime' && vite",
"build": "tsc -b && vite build",
"build2": "vite build",
"lint": "eslint .",
"preview": "vite preview"
},
"dependencies": {
"@devnomic/marquee": "^1.0.3",
"@fortawesome/free-brands-svg-icons": "^6.6.0",
"@headlessui/react": "^2.1.3",
"@preact/signals-react": "^2.2.0",
"@tauri-apps/api": "^2.2.0",
"@techstark/opencv-js": "^4.10.0-release.1",
"@zip.js/zip.js": "^2.7.52",
"konva": "^9.3.14",
"onnxruntime-web": "^1.19.2",
"posthog-js": "^1.174.3",
"react": "^18.3.1",
"react-colorful": "^5.6.1",
"react-dom": "^18.3.1",
"react-drag-drop-files": "^2.3.10",
"react-easy-crop": "^5.1.0",
"react-router-dom": "^6.26.1",
"sonner": "^2.0.1",
"tailwind-merge": "^2.5.2",
"uuid": "^10.0.0",
"yarn": "^1.22.22"
},
"devDependencies": {
"@eslint/js": "^9.9.0",
"@fortawesome/fontawesome-svg-core": "^6.6.0",
"@fortawesome/pro-solid-svg-icons": "^6.6.0",
"@fortawesome/react-fontawesome": "^0.2.2",
"@preact/signals-react-transform": "^0.4.0",
"@tauri-apps/cli": "^2.4.0",
"@types/deep-equal": "^1.0.4",
"@types/react": "^18.3.3",
"@types/react-dom": "^18.3.0",
"@types/uuid": "^10.0.0",
"@vitejs/plugin-react": "^4.3.1",
"autoprefixer": "^10.4.20",
"deep-equal": "^2.2.3",
"eslint": "^9.9.0",
"eslint-plugin-react": "^7.35.0",
"eslint-plugin-react-hooks": "^5.1.0-rc.0",
"eslint-plugin-react-refresh": "^0.4.9",
"globals": "^15.9.0",
"i": "^0.3.7",
"npm": "^10.8.3",
"postcss": "^8.4.41",
"prettier": "^3.3.3",
"prettier-plugin-tailwindcss": "^0.6.6",
"tailwindcss": "^3.4.10",
"typescript": "^5.5.3",
"typescript-eslint": "^8.0.1",
"vite": "^5.4.1",
"vite-plugin-static-copy": "^1.0.6",
"vite-tsconfig-paths": "^5.0.1"
}
}
@@ -0,0 +1,272 @@
import asyncio
import base64
import json
import logging
import pathlib
import time
from io import BytesIO
from typing import Dict, Any, Optional, Tuple, Callable, List
from dataclasses import dataclass
from dataclasses_json import dataclass_json
from enum import Enum, auto
import base64
from PIL import Image
import io
import torch
from diffusers import AutoPipelineForImage2Image, LCMScheduler
from diffusers.utils import make_image_grid, load_image
from sd_embed.embedding_funcs import get_weighted_text_embeddings_sdxl
import torch
import websockets
from PIL import Image
from diffusers import (
StableDiffusionXLPipeline,
StableDiffusionXLImg2ImgPipeline,
AutoPipelineForImage2Image,
LCMScheduler,
UNet2DConditionModel,
)
import torch
import torchvision.transforms as transforms
import gc
transform = transforms.Compose([
transforms.ToTensor(), # This converts PIL Image to tensor and scales to [0, 1]
])
@dataclass_json
@dataclass
class ProgressMessage:
message: str
progress: float
type: str = "progress"
error: Optional[str] = None
@dataclass_json
@dataclass
class LoadModelRequest:
model_path: str
lora_path: Optional[str] = None
command: str = "load_model"
@dataclass_json
@dataclass
class GenerateRequest:
image: str
prompt: str
command: str = "generate"
strength: float = 0.6
guidance_scale: float = 2.0
num_inference_steps: int = 4
height: int = 1024
width: int = 1024
lora_strength: float = 1.0
@dataclass_json
@dataclass
class GenerateResponse:
image: str
type: str = "result"
def save_image(image: Image.Image, base_path: str = "outputs") -> str:
"""Save a PIL Image to disk with a timestamp-based filename."""
try:
# Create outputs directory if it doesn't exist
output_dir = pathlib.Path(base_path)
output_dir.mkdir(parents=True, exist_ok=True)
# Generate filename with timestamp
timestamp = time.strftime("%Y%m%d-%H%M%S")
filename = f"generated-{timestamp}.png"
filepath = output_dir / filename
# Save the image
print(f"Saving image to: {filepath}")
image.save(str(filepath), format="PNG")
print(f"Successfully saved image to: {filepath}")
return str(filepath)
except Exception as e:
print(f"Error saving image: {e}")
return ""
async def send_progress(websocket, message: str, progress: float = 0.0, error: str = None):
"""Send progress updates to the client."""
progress_msg = ProgressMessage(
message=message,
progress=progress,
error=error
)
await websocket.send(progress_msg.to_json())
async def load_model(model_path: str, lora_path: str = None) -> StableDiffusionXLImg2ImgPipeline:
"""Load the model with progress reporting."""
try:
# Load base SDXL pipeline
pipe = StableDiffusionXLPipeline.from_single_file(
model_path,
torch_dtype=torch.float16,
use_safetensors=True
).to("cuda")
# Convert to img2img pipeline
pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
# Set scheduler
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
# Load LCM-LoRA weights if provided
if lora_path:
pipe.load_lora_weights(
"F:/ComfyUI_windows_portable_nvidia/ComfyUI_windows_portable/ComfyUI/models/loras/LCM_LoRA_Weights_SDXL.safetensors",
adapter_name="lcm"
)
pipe.load_lora_weights(lora_path, adapter_name="lora")
pipe.set_adapters(["lcm", "lora"], adapter_weights=[1.0, 1.0])
#pipe.set_adapters(["lcm"], adapter_weights=[1.0])
#pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
return pipe
except Exception as e:
raise Exception(f"Error loading model: {str(e)}")
#https://huggingface.co/blog/simple_sdxl_optimizations
async def generate_image(
pipe: StableDiffusionXLImg2ImgPipeline,
request: GenerateRequest
) -> str:
"""Generate image and return base64 string."""
try:
# Clear CUDA cache before processing
torch.cuda.empty_cache()
image_stream = io.BytesIO(base64.b64decode(request.image))
init_image = Image.open(image_stream).convert('RGB')
init_image = init_image.resize((request.width, request.height)) # width, height
init_image = transform(init_image).to('cuda')
init_image = init_image.unsqueeze(0)
(
prompt_embeds
, prompt_neg_embeds
, pooled_prompt_embeds
, negative_pooled_prompt_embeds
) = get_weighted_text_embeddings_sdxl(
pipe,
prompt = request.prompt,
neg_prompt = "bad quality, poorly Rendered face, poorly drawn face, poorly facial details, poorly drawn hands, poorly rendered hands, low resolution, Images cut out at the top, left, right, bottom, bad composition, mutated body parts, blurry image, disfigured, oversaturated, bad anatomy, deformed body features"
)
# Generate image
with torch.inference_mode():
output_image = pipe(
prompt_embeds = prompt_embeds,
negative_prompt_embeds = prompt_neg_embeds,
pooled_prompt_embeds = pooled_prompt_embeds,
negative_pooled_prompt_embeds = negative_pooled_prompt_embeds,
image=init_image,
num_inference_steps=request.num_inference_steps,
guidance_scale=request.guidance_scale,
strength=request.strength,
generator=torch.manual_seed(0),
height=request.height, width=request.width,
).images[0]
# Clean up CUDA tensors
del init_image, prompt_embeds, prompt_neg_embeds,pooled_prompt_embeds, negative_pooled_prompt_embeds
torch.cuda.empty_cache()
gc.collect()
# Convert to base64
buffered = BytesIO()
output_image.save(buffered, format="PNG")
buffered.seek(0)
img_str = base64.b64encode(buffered.getvalue()).decode()
return f"data:image/png;base64,{img_str}"
except Exception as e:
raise Exception(f"Error generating image: {str(e)} at line {e.__traceback__.tb_lineno}")
async def handle_client(websocket):
"""Handle WebSocket client connection."""
pipe = None
try:
async for message in websocket:
try:
# Clear memory before processing new request
torch.cuda.empty_cache()
gc.collect()
data = json.loads(message)
command = data.get("command")
if command == "load_model":
await send_progress(websocket, "Loading model...", 0.0)
request = LoadModelRequest.from_json(message)
if pipe is not None:
# Clean up existing model
del pipe
torch.cuda.empty_cache()
gc.collect()
pipe = await load_model(request.model_path, request.lora_path)
await send_progress(websocket, "Model loaded successfully", 1.0)
elif command == "generate":
if not pipe:
raise Exception("Model not loaded. Please load model first.")
await send_progress(websocket, "Generating image...", 0.0)
request = GenerateRequest.from_json(message)
# Use a timeout for generation to prevent hanging
try:
result = await asyncio.wait_for(
generate_image(pipe=pipe, request=request),
timeout=10.0 # 30 second timeout
)
response = GenerateResponse(image=result)
await websocket.send(response.to_json())
except asyncio.TimeoutError:
# Force cleanup if generation times out
torch.cuda.empty_cache()
gc.collect()
raise Exception("Generation timed out")
else:
raise Exception(f"Unknown command: {command}")
except Exception as e:
await send_progress(websocket, "Error processing request", 0.0, str(e))
# Clean up after error
torch.cuda.empty_cache()
gc.collect()
except websockets.exceptions.ConnectionClosed:
print("Client connection closed")
except Exception as e:
await send_progress(websocket, "Fatal error", 0.0, str(e))
finally:
# Clean up resources when client disconnects
if pipe is not None:
del pipe
torch.cuda.empty_cache()
gc.collect()
async def server():
"""Start WebSocket server."""
async with websockets.serve(
handle_client,
"localhost",
8765,
max_size=None,
ping_interval=None, # Disable ping/pong to prevent timeouts
max_queue=10 # Limit pending connections
):
print("Server started on ws://localhost:8765")
await asyncio.Future() # run forever
if __name__ == "__main__":
# test() # Comment out test function
asyncio.run(server())
@@ -0,0 +1,101 @@
import asyncio
import base64
import json
import logging
import pathlib
import time
from io import BytesIO
from typing import Dict, Any, Optional, Tuple, Callable, List
import torch
from diffusers import AutoPipelineForImage2Image, LCMScheduler
from diffusers.utils import make_image_grid, load_image
import torch
import websockets
from PIL import Image
from diffusers import (
StableDiffusionXLPipeline,
StableDiffusionXLImg2ImgPipeline,
AutoPipelineForImage2Image,
LCMScheduler,
UNet2DConditionModel,
)
def save_image(image: Image.Image, base_path: str = "outputs") -> str:
"""Save a PIL Image to disk with a timestamp-based filename."""
try:
# Create outputs directory if it doesn't exist
output_dir = pathlib.Path(base_path)
output_dir.mkdir(parents=True, exist_ok=True)
# Generate filename with timestamp
timestamp = time.strftime("%Y%m%d-%H%M%S")
filename = f"generated-{timestamp}.png"
filepath = output_dir / filename
# Save the image
print(f"Saving image to: {filepath}")
image.save(str(filepath), format="PNG")
print(f"Successfully saved image to: {filepath}")
return str(filepath)
except Exception as e:
print(f"Error saving image: {e}")
return ""
def test():
# Load base SDXL pipeline from .safetensors file
pipe = StableDiffusionXLPipeline.from_single_file(
f"C:/Users/Tensor/Downloads/animagineXL40_v4Opt.safetensors",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
# Convert to img2img pipeline
pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
# Set scheduler
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
# Load LCM-LoRA weights
pipe.load_lora_weights(
"F:/ComfyUI_windows_portable_nvidia/ComfyUI_windows_portable/ComfyUI/models/loras/LCM_LoRA_Weights_SDXL.safetensors",
adapter_name="lcm"
)
pipe.set_adapters(["lcm"], adapter_weights=[1.0])
# prepare image
url = "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/7462a0ac-637d-467b-baa9-34c5f54696b4/original=true,quality=90/60277788.jpeg"
init_image = load_image(url)
prompt = "safe_pos, 1girl, solo, solo focus, 8k, masterpiece, hires, absurdres, splash art, gradient, looking at viewer, medium breasts, vibrant yellow hair, purple hair tips, (shoulder freckles:1.2), nude, athletic, confident, hand on waist, extremely long hair, dancing, delicate gold jewelry, thin gold bangles, belly button chain, action shot,, head tilt, hair censor, spotlight, tottotonero"
# pass prompt and image to pipeline
generator = torch.manual_seed(0)
output_image = pipe(
prompt,
image=init_image,
num_inference_steps=4,
guidance_scale=2,
strength=0.6,
generator=generator
).images[0]
# Save both input and output images
save_image(init_image, "outputs/input")
output_path = save_image(output_image, "outputs/output")
print(f"Output image saved to: {output_path}")
# Create and save the grid
grid = make_image_grid([init_image, output_image], rows=1, cols=2)
grid_path = save_image(grid, "outputs/grid")
print(f"Grid image saved to: {grid_path}")
async def start_server():
async with websockets.serve(
"localhost",
8765 # Increase to 10MB
):
await asyncio.Future() # run forever
if __name__ == "__main__":
test()
+6
View File
@@ -0,0 +1,6 @@
export default {
plugins: {
tailwindcss: {},
autoprefixer: {},
},
}

Before

Width:  |  Height:  |  Size: 16 KiB

After

Width:  |  Height:  |  Size: 16 KiB

Some files were not shown because too many files have changed in this diff Show More