Compare commits

...
Author SHA1 Message Date
Brandon Thomas 14feaccb91 Merge branch 'main' into feature/generatemode 2025-03-25 01:38:48 -04:00
Brandon Thomas 81c256e4f8 cleanup 2025-03-24 23:08:07 -04:00
Brandon Thomas 8d51d5b8ab can't trigger default features; switch to platform switches 2025-03-24 23:00:29 -04:00
Brandon Thomas 456bfd7bb9 trigger with --feature ml_models (need to make default) 2025-03-24 22:31:58 -04:00
Brandon Thomas cd3c864eff runs on mac 2025-03-24 21:46:25 -04:00
Brandon Thomas 990b7ff3e5 package split for mac (part one) (#459) 2025-03-24 20:18:57 -04:00
Brandon Thomas e1dbf4f2fe add inpaint demo/test endpoint 2025-03-22 03:16:48 -04:00
Tanish Manku d826db599d Remove the cancel generation button for now 2025-03-22 00:08:06 -07:00
Tanish Manku a642732c59 Clip overflow to have top rounded corners on image 2025-03-22 00:02:21 -07:00
Brandon Thomas 435c5aa0b4 fix strength bug 2025-03-22 02:57:06 -04:00
Tanish Manku 90b4aff846 Fix: Multiple invokes of a prompt after changing mode 2025-03-21 23:56:50 -07:00
Tanish Manku 4d3f097432 Hook up the TTI UI to invoke call 2025-03-21 23:38:25 -07:00
Tanish Manku e5de0e4731 Merge branch 'main' into feature/generatemode 2025-03-21 23:11:07 -07:00
Tanish Manku 04498b0457 Add initial Edit mode files 2025-03-21 23:02:17 -07:00
bflatastic fbfc1e15f2 gallery gap fix 2025-03-21 22:46:30 +07:00
bflatastic 4a65317729 rename image to realtime 2025-03-21 20:13:13 +07:00
Tanish Manku e7ed0693a3 Add generation engine and signals 2025-03-21 00:44:44 -07:00
Tanish Manku 99ac8faf33 Merge branch 'main' into feature/generatemode 2025-03-20 23:41:19 -07:00
bflatastic f0213d4a63 cancel button 2025-03-21 00:03:15 +07:00
bflatastic 58f684b1e0 tab disabled prop and disabled message to block tab change when generating 2025-03-20 22:42:36 +07:00
bflatastic e83389a874 background gallery + many other styling and animation transition updates 2025-03-20 22:11:12 +07:00
Tanish Manku 3602f8fc2b Add generation signals 2025-03-20 02:10:27 -07:00
Tanish Manku 115c4fc60c Add a generation mode UI component 2025-03-19 23:58:53 -07:00
91 changed files with 1713 additions and 439 deletions
Generated
+54 -9
View File
@@ -661,25 +661,19 @@ dependencies = [
"assertor",
"base64 0.22.1",
"bytes 1.10.0",
"candle-core",
"candle-nn",
"candle-transformers",
"const_format",
"cudarc",
"directories 6.0.0",
"dirs 6.0.0",
"filesys",
"futures 0.3.29",
"half",
"hf-hub 0.3.2",
"image",
"libloading 0.8.6",
"log",
"memory_caching",
"ndarray",
"nvidia-checker",
"ml_models",
"ml_weights_registry",
"once_cell",
"ort",
"primitives",
"rand 0.8.5",
"reqwest 0.12.12",
@@ -692,7 +686,6 @@ dependencies = [
"tauri-plugin-log",
"tempdir",
"tempfile",
"tokenizers 0.15.2",
"tokio",
"tokio-util",
]
@@ -6804,6 +6797,58 @@ dependencies = [
"ws2_32-sys",
]
[[package]]
name = "ml_models"
version = "0.0.1"
dependencies = [
"anyhow",
"assertor",
"base64 0.22.1",
"bytes 1.10.0",
"candle-core",
"candle-nn",
"candle-transformers",
"const_format",
"cudarc",
"directories 6.0.0",
"dirs 6.0.0",
"filesys",
"futures 0.3.29",
"half",
"hf-hub 0.3.2",
"image",
"libloading 0.8.6",
"log",
"memory_caching",
"ml_weights_registry",
"ndarray",
"nvidia-checker",
"once_cell",
"ort",
"primitives",
"rand 0.8.5",
"reqwest 0.12.12",
"serde",
"serde_json",
"serde_yaml",
"serial_test",
"tauri",
"tauri-plugin-log",
"tempdir",
"tempfile",
"tokenizers 0.15.2",
"tokio",
"tokio-util",
]
[[package]]
name = "ml_weights_registry"
version = "0.0.1"
dependencies = [
"anyhow",
"const_format",
]
[[package]]
name = "mockall"
version = "0.12.1"
+2
View File
@@ -23,6 +23,8 @@ members = [
# Desktop apps
"crates/desktop/src-tauri",
"crates/desktop/ml_models",
"crates/desktop/ml_weights_registry",
# Shared libraries
"crates/lib/actix_cors_configs",
+79
View File
@@ -0,0 +1,79 @@
[package]
name = "ml_models"
edition = "2021"
version = "0.0.1"
authors = ["storyteller"]
publish = false
description = ""
license = ""
repository = ""
[lib]
name = "ml_models"
path = "src/lib.rs"
crate-type = ["staticlib", "cdylib", "rlib"]
[package.metadata.cargo-udeps.ignore]
normal = ["my-workspace-hack"]
[dependencies]
# Internal
filesys = { path = "../../lib/files/filesys" }
memory_caching = { path = "../../lib/caching/memory_caching" }
ml_weights_registry = { path = "../ml_weights_registry" }
primitives = { path = "../../lib/data/primitives" }
# Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332
cudarc = "0.13.7"
# External
anyhow = "1.0"
base64 = "0.22.1"
bytes = "1.10"
candle-core = { path = "../../vendor/candle/candle-core", features = ["cuda"] }
candle-nn = { path = "../../vendor/candle/candle-nn", features = ["cuda"] }
candle-transformers = { path = "../../vendor/candle/candle-transformers", features = ["cuda"] }
const_format = "0.2.34"
directories = "6.0.0"
dirs = "6.0.0"
futures = "0.3"
half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
hf-hub = "0.3.2"
image = "0.25.5"
log = "0.4"
ndarray = "0.15.6"
nvidia-checker = "0.1.2"
once_cell = "1.20.3"
ort = { version = "1.16.2", features = ["cuda"] }
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"
serde_yaml = "0.9"
tauri = { version = "2.2.4", features = [] }
tauri-plugin-log = "2"
tempdir = "0.3.7"
tempfile = "3.18.0"
tokenizers = "0.15.2"
tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] }
tokio-util = { version = "0.7", features = ["codec"] }
# tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] }
[features]
default = ["cuda"]
cuda = [
"candle-core/cuda",
"candle-nn/cuda",
"candle-transformers/cuda",
# NB(bt,2025-03-10): Building without Flash attention on Windows for now to bypass linker issues
"candle-transformers/flash-attn",
"candle-core/cudnn"
]
[dev-dependencies]
libloading = "0.8.6"
serial_test = "3.2.0"
assertor = "0.0.3"
+1
View File
@@ -0,0 +1 @@
pub mod ml;
@@ -0,0 +1,2 @@
pub mod onnx_session;
pub mod remove_image_background;
@@ -1,8 +1,9 @@
//! Adapted from https://github.com/dnanhkhoa/rust-background-removal/blob/main/src/onnx.rs
use std::path::Path;
use ort::{Environment, ExecutionProvider, GraphOptimizationLevel, SessionBuilder};
pub fn onnx_session(onnx_model_file: &str) -> Result<ort::Session, anyhow::Error> {
pub fn onnx_session(onnx_model_file: &Path) -> Result<ort::Session, anyhow::Error> {
let environment = Environment::default().into_arc();
let session = SessionBuilder::new(&environment)?
.with_optimization_level(GraphOptimizationLevel::Level1)? // Configure model optimization level
@@ -0,0 +1,61 @@
use crate::ml::background_removal::onnx_session::onnx_session;
use image::{imageops, DynamicImage};
use ndarray::{Array, CowArray};
use ort::Value;
use std::path::Path;
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:
// sudo apt-get install libcudnn8-dev libcudnn8
// update-alternatives: warning: forcing reinstallation of alternative /usr/include/x86_64-linux-gnu/cudnn_v9.h because link group libcudnn is broken
// 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)?;
Ok(image)
}
fn process_dynamic_image(
session: &ort::Session,
dynamic_img: DynamicImage,
) -> Result<DynamicImage, anyhow::Error> {
let input_shape = session.inputs[0]
.dimensions()
.map(|dim| dim.unwrap())
.collect::<Vec<usize>>();
let input_img = dynamic_img.into_rgba8();
let scaling_factor = f32::min(
1., // Avoid upscaling
f32::min(
input_shape[3] as f32 / input_img.width() as f32, // Width ratio
input_shape[2] as f32 / input_img.height() as f32, // Height ratio
),
);
let mut resized_img = imageops::resize(
&input_img,
input_shape[3] as u32,
input_shape[2] as u32,
imageops::FilterType::Triangle,
);
let input_tensor = CowArray::from(
Array::from_shape_fn(input_shape, |indices| {
let mean = 128.;
let std = 256.;
(resized_img[(indices[3] as u32, indices[2] as u32)][indices[1]] as f32 - mean) / std
})
.into_dyn(),
);
let inputs = vec![Value::from_array(session.allocator(), &input_tensor)?];
let outputs = session.run(inputs)?;
let output_tensor = outputs[0].try_extract::<f32>()?;
for (indices, alpha) in output_tensor.view().indexed_iter() {
resized_img[(indices[3] as u32, indices[2] as u32)][3] = (alpha * 255.) as u8;
}
let output_img = imageops::resize(
&resized_img,
(input_img.width() as f32 * scaling_factor) as u32,
(input_img.height() as f32 * scaling_factor) as u32,
imageops::FilterType::Triangle,
);
Ok(DynamicImage::ImageRgba8(output_img))
}
@@ -1,7 +1,7 @@
pub mod background_removal;
pub mod image;
pub mod model_cache;
pub mod model_config;
pub mod model_file;
pub mod prompt_cache;
pub mod stable_diffusion;
pub mod weights_registry;
pub mod stable_diffusion;
@@ -0,0 +1,61 @@
use crate::ml::model_file::StableDiffusionVersion;
use candle_core::{DType, Device};
use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
use hf_hub::api::sync::Api;
use log::info;
const DEFAULT_SD_IMAGE_WIDTH: usize = 1024;
const DEFAULT_SD_IMAGE_HEIGHT: usize = 1024;
const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32;
const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::V1_5;
pub struct ModelConfig {
/// We only support one device type currently, and we can't multiplex
/// over several devices yet.
pub device: Device,
/// The core numeric type to use for models.
pub dtype: DType,
pub sd_version: StableDiffusionVersion,
pub sd_config: StableDiffusionConfig,
/// Probably shouldn't be here.
pub hf_api: Api,
}
impl ModelConfig {
pub fn init() -> anyhow::Result<Self> {
let device = Device::new_cuda(0)
.unwrap_or_else(|e| {
info!("CUDA not available ({}), falling back to CPU", e);
Device::Cpu
});
let sd_config = match SD_VERSION {
StableDiffusionVersion::Turbo => {
StableDiffusionConfig::sdxl_turbo(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
}
StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => {
StableDiffusionConfig::v1_5(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
}
StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint => {
StableDiffusionConfig::v2_1(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
}
StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint => {
StableDiffusionConfig::sdxl(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
}
};
let hf_api = Api::new()?;
Ok(Self {
device,
dtype: DEFAULT_SD_NUMERIC_DATATYPE,
sd_version: SD_VERSION,
sd_config,
hf_api,
})
}
}
@@ -4,18 +4,50 @@ extern crate accelerate_src;
extern crate intel_mkl_src;
use candle_transformers::models::stable_diffusion;
use std::path::PathBuf;
use std::path::{Path, PathBuf};
use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
use crate::ml::weights_registry::weights::{CLIP_JSON, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2};
use crate::state::app_dir::AppWeightsDir;
use ml_weights_registry::weights_registry::weights::{CLIP_JSON, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2};
use anyhow::Error as E;
use candle_core::{DType, Device, Module, Tensor, D};
use log::info;
use tokenizers::Tokenizer;
pub struct ClipArgs<'a, P: AsRef<Path>, Q: AsRef<Path>> {
pub prompt: &'a str,
pub uncond_prompt: &'a str,
pub tokenizer: Option<String>,
pub clip_weights: Option<String>,
pub clip2_weights: Option<String>,
pub sd_version: StableDiffusionVersion,
pub sd_config: &'a stable_diffusion::StableDiffusionConfig,
pub use_f16: bool,
pub device: &'a Device,
pub dtype: DType,
pub use_guide_scale: bool,
pub clip_json_path: P,
pub clip_weights_path: Q,
}
#[allow(clippy::too_many_arguments)]
pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: Option<String>, clip_weights: Option<String>, clip2_weights: Option<String>, sd_version: StableDiffusionVersion, sd_config: &stable_diffusion::StableDiffusionConfig, use_f16: bool, device: &Device, dtype: DType, use_guide_scale: bool, weights_dir: &AppWeightsDir) -> anyhow::Result<Tensor> {
pub fn infer_clip_text_embeddings<P: AsRef<Path>, Q: AsRef<Path>>(args: ClipArgs<'_, P, Q>) -> anyhow::Result<Tensor> {
let ClipArgs {
prompt,
uncond_prompt,
tokenizer,
clip_weights,
clip2_weights,
sd_version,
sd_config,
use_f16,
device,
dtype,
use_guide_scale,
clip_json_path,
clip_weights_path,
} = args;
info!("Preparing text embeddings... for {:?}", sd_version);
let which = match sd_version {
@@ -39,7 +71,8 @@ pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer:
dtype,
use_guide_scale,
*first,
weights_dir,
clip_json_path.as_ref(),
clip_weights_path.as_ref(),
)
})
.collect::<anyhow::Result<Vec<_>>>()?;
@@ -50,10 +83,25 @@ pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer:
}
#[allow(clippy::too_many_arguments)]
fn do_infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: Option<String>, clip_weights: Option<String>, clip2_weights: Option<String>, sd_version: StableDiffusionVersion, sd_config: &stable_diffusion::StableDiffusionConfig, use_f16: bool, device: &Device, dtype: DType, use_guide_scale: bool, first: bool, weights_dir: &AppWeightsDir) -> anyhow::Result<Tensor> {
fn do_infer_clip_text_embeddings(
prompt: &str,
uncond_prompt: &str,
tokenizer: Option<String>,
clip_weights: Option<String>,
clip2_weights: Option<String>,
sd_version: StableDiffusionVersion,
sd_config: &stable_diffusion::StableDiffusionConfig,
use_f16: bool,
device: &Device,
dtype: DType,
use_guide_scale: bool,
first: bool,
clip_json_path: &Path,
clip_weights_path: &Path,
) -> anyhow::Result<Tensor> {
info!("do_infer_clip_text_embeddings called with args {:?}", (prompt, uncond_prompt, tokenizer.clone().unwrap_or_else(|| "None".to_string()), clip_weights.clone().unwrap_or_else(|| "None".to_string()), clip2_weights.clone().unwrap_or_else(|| "None".to_string()), sd_version, sd_config, use_f16, device, dtype, use_guide_scale, first));
let tokenizer = weights_dir.weight_path(&CLIP_JSON);
let tokenizer = clip_json_path;
info!("Loading Clip Tokenizer path: {:?}", tokenizer);
@@ -75,11 +123,16 @@ fn do_infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: O
let tokens = Tensor::new(tokens.as_slice(), device)?.unsqueeze(0)?;
println!("Building the Clip transformer.");
// TODO(bt,2025-03-23): These used to be two different clip models. We should only need one.
let clip_weights = if first {
weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER)
println!(">>>>> CLIP 1");
clip_weights_path
} else {
weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER_2)
println!(">>>>> CLIP 2");
clip_weights_path
};
let clip_config = if first { &sd_config.clip } else { sd_config.clip2.as_ref().unwrap() };
let text_model = stable_diffusion::build_clip_transformer(clip_config, clip_weights, device, DType::F32)?;
@@ -1,53 +1,80 @@
use std::path::PathBuf;
use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
use crate::events::notification_event::{NotificationEvent, NotificationModelType};
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::model_file::{ModelFile, StableDiffusionVersion};
use crate::ml::prompt_cache::PromptCache;
use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_embeddings;
use crate::ml::stable_diffusion::remap_lcm_strength_range::remap_lcm_strength_range;
use crate::ml::weights_registry::weights::{LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
use crate::state::app_config::AppConfig;
use crate::state::app_dir::AppDataRoot;
use anyhow::{anyhow, Error as E, Result};
use candle_core::{DType, IndexOp, Tensor, D};
use candle_core::{DType, Device, IndexOp, Tensor, D};
use candle_nn::VarBuilder;
use candle_transformers::models::stable_diffusion::lcm::LCMScheduler;
use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
use candle_transformers::models::stable_diffusion::lcm::{LCMScheduler, LCMSchedulerConfig};
use candle_transformers::models::stable_diffusion::unet_2d::BlockConfig;
use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModelConfig;
use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
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::model_config::ModelConfig;
use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
use crate::ml::prompt_cache::PromptCache;
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 ml_weights_registry::weights_registry::weights::{LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
use rand::Rng;
use tauri::{AppHandle, Emitter};
const DEFAULT_STRENGTH : f64 = 75.0;
pub struct Args<'a> {
pub struct Args<'a, P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>> {
pub image: &'a DynamicImage,
pub prompt: String,
pub uncond_prompt: String,
pub cfg_scale: Option<f64>,
pub i2i_strength: Option<f64>,
pub configs: &'a AppConfig,
pub img2img_strength: Option<f64>,
pub model_cache: &'a ModelCache,
pub prompt_cache: &'a PromptCache,
pub app_data_root: &'a AppDataRoot,
pub app: &'a AppHandle,
pub use_flash_attn: bool,
pub model_config: &'a ModelConfig,
pub maybe_seed: Option<u64>,
pub scheduler_steps: usize,
pub vae_path: P1,
pub unet_path: P2,
pub clip_json_path: P3,
pub clip_weights_path: P4,
}
pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
let Args { prompt, uncond_prompt, cfg_scale, i2i_strength, configs, model_cache, prompt_cache, app, image, app_data_root, use_flash_attn } = args;
let weights_dir = app_data_root.weights_dir();
pub fn lcm_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<RgbImage>
where P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>
{
let Args {
prompt,
uncond_prompt,
cfg_scale,
img2img_strength,
model_cache,
prompt_cache,
image,
use_flash_attn ,
model_config: ModelConfig {
device,
dtype,
sd_version,
sd_config,
hf_api,
},
maybe_seed,
scheduler_steps,
vae_path,
unet_path,
clip_json_path,
clip_weights_path,
} = args;
let img2img_strength = args.i2i_strength
let dtype = *dtype;
let sd_version = *sd_version;
let img2img_strength = img2img_strength
.map(|s| remap_lcm_strength_range(s))
.unwrap_or(DEFAULT_STRENGTH);
@@ -56,14 +83,17 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
let use_f16 = true;
// Use LCM Scheduler instead of Euler Ancestral for better speed and quality
let mut scheduler = LCMScheduler::new(configs.scheduler_steps, img2img_strength, candle_transformers::models::stable_diffusion::lcm::LCMSchedulerConfig::default())?;
let mut scheduler = LCMScheduler::new(scheduler_steps, img2img_strength, LCMSchedulerConfig::default())?;
let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
configs.device.set_seed(seed)?;
let seed = maybe_seed.unwrap_or_else(|| rand::thread_rng().gen());
let seed = 42;
device.set_seed(seed)?;
let guidance_scale = match cfg_scale {
Some(guidance_scale) => guidance_scale,
None => match configs.sd_version {
None => match sd_version {
StableDiffusionVersion::V1_5 | StableDiffusionVersion::V2_1 | StableDiffusionVersion::Xl => 7.5,
StableDiffusionVersion::Turbo => 0.,
_ => 0., // NB: Not sure what the other model families should use, so sticking with "0"
@@ -78,50 +108,45 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
let mut text_embeddings = if let Some(tensor) = maybe_cached {
tensor
} else {
info!("Loading clip weights...");
let clip_weights = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16);
let clip_weights = Some(clip_weights.to_string_lossy().to_string());
info!("Prompt is NOT cached! Calculating embedding...");
let tensor = infer_clip_text_embeddings(
&prompt,
&uncond_prompt,
None, // tokenizer
clip_weights,
None, // clip2_weights
configs.sd_version,
&configs.sd_config,
let tensor = infer_clip_text_embeddings(ClipArgs {
prompt: & prompt,
uncond_prompt: &uncond_prompt,
tokenizer: None, // tokenizer
clip_weights: None,
clip2_weights: None, // clip2_weights
sd_version,
sd_config: &sd_config,
use_f16, // use_f16
&configs.device,
configs.dtype,
device,
dtype,
use_guide_scale,
weights_dir,
)?;
clip_json_path,
clip_weights_path,
})?;
prompt_cache.store_copy(&prompt, &tensor)?;
tensor
};
println!("Text embeddings shape: {:?}", text_embeddings.shape());
info!("Text embeddings shape: {:?}", text_embeddings.shape());
println!("Loading input image into tensor...");
info!("Loading input image into tensor...");
let input_image = dynamic_image_to_tensor(image, &configs.device, configs.dtype)?;
let input_image = dynamic_image_to_tensor(image, &device, dtype)?;
println!("Reference image shape: {:?}", input_image.shape());
info!("Reference image shape: {:?}", input_image.shape());
let vae_scale = get_vae_scale(configs.sd_version);
let vae_scale = get_vae_scale(sd_version);
let maybe_vae = model_cache.get_vae()?;
let vae = match maybe_vae {
Some(vae) => vae,
None => {
let vae_file = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE);
println!("Building VAE model from file {:?}...", &vae_file);
let vae = configs.sd_config.build_vae(vae_file, &configs.device, configs.dtype)?;
info!("Building VAE model from file {:?}...", vae_path.as_ref());
let vae = sd_config.build_vae(vae_path, &device, dtype)?;
let vae = Arc::new(vae);
@@ -138,9 +163,7 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
None => {
info!("No unet found in cache; loading...");
let unet_weights = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
let in_channels = match configs.sd_version {
let in_channels = match sd_version {
StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::V1_5Inpaint => 9,
_ => 4,
};
@@ -175,7 +198,7 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
};
// Load the model directly from the safetensors file
let vs_unet = unsafe { VarBuilder::from_mmaped_safetensors(&[unet_weights], configs.dtype, &configs.device)? };
let vs_unet = unsafe { VarBuilder::from_mmaped_safetensors(&[unet_path], dtype, &device)? };
let unet = UNet2DConditionModel::new(vs_unet, in_channels, 4, use_flash_attn, unet_config)?;
@@ -191,19 +214,19 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
// Create LCM guidance scale embeddings
let embedding_dim = 1280;
let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &configs.device, configs.dtype)?;
let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &device, dtype)?;
// Generate latents from input image using the LCM approach
println!("Generating latents from input image...");
info!("Generating latents from input image...");
let init_latents = (init_latent_dist.sample()? * vae_scale)?;
println!("Initial latents shape: {:?}", init_latents.shape());
info!("Initial latents shape: {:?}", init_latents.shape());
// Calculate timesteps for LCM with proper img2img handling
let timesteps = LCMScheduler::get_timesteps_for_steps(configs.scheduler_steps, img2img_strength);
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
println!("Adding noise to latents for t_start={}", t_start);
info!("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)?;
@@ -211,7 +234,7 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect();
for (timestep_index, &timestep) in timesteps.iter().enumerate() {
println!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep);
info!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep);
let latent_model_input = if use_guide_scale { Tensor::cat(&[&latents, &latents], 0)? } else { latents.clone() };
@@ -221,7 +244,7 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
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) => {
println!("UNet inference failed with error: {}", e);
info!("UNet inference failed with error: {}", e);
return Err(anyhow::anyhow!("UNet inference failed: {}", e));
},
};
@@ -233,24 +256,24 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
}
// Apply the LCM scheduler step
latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, configs.scheduler_steps)?;
latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, scheduler_steps)?;
}
println!("Diffusion process completed, decoding image...");
info!("Diffusion process completed, decoding image...");
let image = vae.decode(&(latents / vae_scale)?)?;
println!("VAE decode completed");
info!("VAE decode completed");
println!("Scaling image back");
info!("Scaling image back");
let image = ((image / 2.)? + 0.5)?;
println!("Normalized image values");
info!("Normalized image values");
let image = (image.clamp(0f32, 1.)? * 255.)?;
println!("Convert to int8");
info!("Convert to int8");
let image = image.to_dtype(DType::U8)?;
println!("Converted to 8-bit format");
info!("Converted to 8-bit format");
let image = tensor_to_image_buffer(&image.i(0)?)?;
@@ -1,75 +1,86 @@
use std::path::Path;
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
use crate::events::notification_event::{NotificationEvent, NotificationModelType};
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::model_file::StableDiffusionVersion;
use crate::ml::prompt_cache::PromptCache;
use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_embeddings;
use crate::ml::weights_registry::weights::{SDXL_TURBO_UNET, SDXL_TURBO_VAE};
use crate::state::app_config::AppConfig;
use crate::state::app_dir::AppDataRoot;
use crate::ml::stable_diffusion::infer_clip_text_embeddings::{infer_clip_text_embeddings, ClipArgs};
use anyhow::{anyhow, Error as E, Result};
use candle_core::{DType, IndexOp, Tensor, D};
use candle_core::{DType, Device, IndexOp, Tensor, D};
use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
use image::DynamicImage;
use log::info;
use rand::Rng;
use tauri::{AppHandle, Emitter};
pub struct Args<'a> {
pub struct Args<'a, P: AsRef<Path>, Q: AsRef<Path>, R: AsRef<Path>, S: AsRef<Path>> {
pub image: &'a DynamicImage,
pub prompt: String,
pub uncond_prompt: String,
pub cfg_scale: Option<f64>,
pub i2i_strength: Option<u8>,
pub configs: &'a AppConfig,
pub model_cache: &'a ModelCache,
pub prompt_cache: &'a PromptCache,
pub app_data_root: &'a AppDataRoot,
pub app: &'a AppHandle,
pub sd_version: StableDiffusionVersion,
pub sd_config: StableDiffusionConfig,
pub device: &'a Device,
pub dtype: DType,
pub maybe_seed: Option<u64>,
pub scheduler_steps: usize,
pub sdxl_turbo_vae_path: P,
pub sdxl_turbo_unet_path: Q,
pub clip_json_path: R,
pub clip_weights_path: S,
}
pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
pub fn stable_diffusion_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<RgbImage>
where P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>
{
let Args {
prompt,
uncond_prompt,
cfg_scale,
i2i_strength,
configs,
model_cache,
prompt_cache,
app,
image,
app_data_root,
sd_version,
sd_config,
device,
dtype,
maybe_seed,
scheduler_steps,
sdxl_turbo_vae_path,
sdxl_turbo_unet_path,
clip_json_path,
clip_weights_path,
} = args;
let weights_dir = app_data_root.weights_dir();
println!("Starting image generation with the following configuration:");
println!(" Model: {:?}", configs.sd_version);
println!(" Model: {:?}", sd_version);
println!(" Prompt: {}", prompt);
println!(" Steps: {}", configs.scheduler_steps);
println!(" Device: {:?}", configs.device);
println!(" Steps: {}", scheduler_steps);
println!(" Device: {:?}", device);
println!("Model dimensions: {}x{}", configs.sd_config.width, configs.sd_config.height);
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
let mut scheduler = configs.sd_config.build_scheduler(
configs.scheduler_steps)?;
let mut scheduler = sd_config.build_scheduler(scheduler_steps)?;
let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
configs.device.set_seed(seed)?;
let seed = maybe_seed.unwrap_or_else(|| rand::thread_rng().gen());
device.set_seed(seed)?;
info!("Using seed: {}", seed);
let guidance_scale = match cfg_scale {
Some(guidance_scale) => guidance_scale,
None => match configs.sd_version {
None => match sd_version {
StableDiffusionVersion::V1_5
| StableDiffusionVersion::V2_1
| StableDiffusionVersion::Xl => 7.5,
@@ -90,20 +101,21 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
} else {
info!("Prompt is NOT cached! Calculating embedding...");
let tensor = infer_clip_text_embeddings(
&prompt,
&uncond_prompt,
None, // tokenizer
None, // clip_weights
None, // clip2_weights
configs.sd_version,
&configs.sd_config,
false, // use_f16
&configs.device,
configs.dtype,
let tensor = infer_clip_text_embeddings(ClipArgs {
prompt: &prompt,
uncond_prompt: &uncond_prompt,
tokenizer: None, // tokenizer
clip_weights: None, // clip_weights
clip2_weights: None, // clip2_weights
sd_version: sd_version,
sd_config: &sd_config,
use_f16: false, // use_f16
device: device,
dtype: dtype,
use_guide_scale,
weights_dir,
)?;
clip_json_path: clip_json_path.as_ref(),
clip_weights_path: clip_weights_path.as_ref(),
})?;
prompt_cache.store_copy(&prompt, &tensor)?;
tensor
};
@@ -112,14 +124,11 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
println!("Loading input image into tensor...");
let input_image = dynamic_image_to_tensor(
image,
&configs.device,
configs.dtype)?;
let input_image = dynamic_image_to_tensor(image, device, dtype)?;
println!("Reference image shape: {:?}", input_image.shape());
let vae_scale = get_vae_scale(configs.sd_version);
let vae_scale = get_vae_scale(sd_version);
let maybe_vae = model_cache.get_vae()?;
@@ -128,11 +137,9 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
None => {
info!("No vae found in cache; loading...");
let vae_file = weights_dir.weight_path(&SDXL_TURBO_VAE);
let vae_file = sdxl_turbo_vae_path.as_ref();
let vae = configs
.sd_config
.build_vae(vae_file, &configs.device, configs.dtype)?;
let vae = sd_config.build_vae(vae_file, &device, dtype)?;
let vae = Arc::new(vae);
@@ -149,9 +156,9 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
None => {
info!("No unet found in cache; loading...");
let unet_weights_file = weights_dir.weight_path(&SDXL_TURBO_UNET);
let unet_weights_file = sdxl_turbo_unet_path.as_ref();
let unet = configs.sd_config.build_unet(unet_weights_file, &configs.device, 4, true, configs.dtype)
let unet = sd_config.build_unet(unet_weights_file, &device, 4, true, dtype)
.map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
let unet = Arc::new(unet);
@@ -181,9 +188,9 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
};
let t_start = {
let start = configs.scheduler_steps - (configs.scheduler_steps as f64 * img2img_strength) as usize;
let start = scheduler_steps - (scheduler_steps as f64 * img2img_strength) as usize;
println!("Starting from step {} of {} (strength: {})", start, configs.scheduler_steps, img2img_strength);
println!("Starting from step {} of {} (strength: {})", start, scheduler_steps, img2img_strength);
start
};
@@ -0,0 +1,18 @@
[package]
name = "ml_weights_registry"
edition = "2021"
version = "0.0.1"
authors = ["storyteller"]
publish = false
description = ""
license = ""
repository = ""
[lib]
name = "ml_weights_registry"
path = "src/lib.rs"
# External
[dependencies]
anyhow = "1.0"
const_format = "0.2.34"
@@ -0,0 +1 @@
pub mod weights_registry;
@@ -5,10 +5,10 @@ pub (super) const R2_BUCKET_URL : &str = "https://pub-bc5e2bc0cdee4bb5ae8fca9d64
// and variables (even if they're known to be &'static) and keep allocations static.
macro_rules! weight {
($name: literal, $file_name:literal, $weight_function:expr) => {
$crate::ml::weights_registry::weight_descriptor::WeightDescriptor {
$crate::weights_registry::weight_descriptor::WeightDescriptor {
name: $name,
filename: $file_name,
r2_download_url: const_format::concatcp!(crate::ml::weights_registry::weight_descriptor_builder::R2_BUCKET_URL, $file_name),
r2_download_url: const_format::concatcp!(crate::weights_registry::weight_descriptor_builder::R2_BUCKET_URL, $file_name),
function: $weight_function,
}
};
@@ -1,5 +1,5 @@
use crate::ml::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction};
use crate::ml::weights_registry::weight_descriptor_builder::weight;
use crate::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction};
use crate::weights_registry::weight_descriptor_builder::weight;
pub const CLIP_JSON : WeightDescriptor = weight!(
"CLIP",
+11 -20
View File
@@ -7,6 +7,7 @@ license = ""
repository = ""
edition = "2021"
rust-version = "1.77.2"
publish = false
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -22,31 +23,26 @@ tauri-build = { version = "2.0.4", features = [] }
# Internal
filesys = { path = "../../lib/files/filesys" }
memory_caching = { path = "../../lib/caching/memory_caching" }
ml_weights_registry = { path = "../ml_weights_registry" }
primitives = { path = "../../lib/data/primitives" }
# Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332
cudarc = "0.13.7"
# Internal ML
# TODO: Only on Windows and Linux
#ml_models = { path = "../ml_models", optional = true }
# External
#openssl = { version = "0.10", features = ["vendored"] }
anyhow = "1.0"
base64 = "0.22.1"
bytes = "1.10"
candle-core = { path = "../../vendor/candle/candle-core", features = ["cuda"] }
candle-nn = { path = "../../vendor/candle/candle-nn", features = ["cuda"] }
candle-transformers = { path = "../../vendor/candle/candle-transformers", features = ["cuda"] }
const_format = "0.2.34"
directories = "6.0.0"
dirs = "6.0.0"
futures = "0.3"
half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
hf-hub = "0.3.2"
image = "0.25.5"
log = "0.4"
ndarray = "0.15.6"
nvidia-checker = "0.1.2"
once_cell = "1.20.3"
ort = { version = "1.16.2", features = ["cuda"] }
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 }
@@ -57,28 +53,23 @@ tauri = { version = "2.2.4", features = [] }
tauri-plugin-log = "2"
tempdir = "0.3.7"
tempfile = "3.18.0"
tokenizers = "0.15.2"
tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] }
tokio-util = { version = "0.7", features = ["codec"] }
# tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] }
[features]
default = ["cuda"]
cuda = [
"candle-core/cuda",
"candle-nn/cuda",
"candle-transformers/cuda",
# NB(bt,2025-03-10): Building without Flash attention on Windows for now to bypass linker issues
"candle-transformers/flash-attn",
"candle-core/cudnn"
]
[target.'cfg(not(target_os = "macos"))'.dependencies]
ml_models = { path = "../ml_models" }
[dev-dependencies]
libloading = "0.8.6"
serial_test = "3.2.0"
assertor = "0.0.3"
#[target.'cfg(target_os = "macos")'.features]
#default = []
#ml_models = ["dep:ml_models"]
[profile.dev]
incremental = true # Compile your binary in smaller steps.
@@ -0,0 +1,24 @@
use crate::ml::model_cache::ModelCache;
use crate::ml::prompt_cache::PromptCache;
use crate::state::app_config::AppConfig;
use crate::state::app_dir::AppDataRoot;
use image::{ImageReader, RgbImage};
use log::info;
use tauri::{AppHandle, State};
#[tauri::command]
pub async fn inpaint_image(
image: &str,
mask: &str,
prompt: String,
model_config: State<'_, AppConfig>,
model_cache: State<'_, ModelCache>,
prompt_cache: State<'_, PromptCache>,
app_data_root: State<'_, AppDataRoot>,
app: AppHandle,
) -> Result<String, String> {
info!("inpaint_endpoint endpoint called.");
Ok(mask.to_string())
}
@@ -1,6 +1,6 @@
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 test_counter;
pub mod text_to_image_endpoint;
@@ -1,16 +1,21 @@
use crate::ml::model_cache::ModelCache;
use crate::ml::prompt_cache::PromptCache;
use crate::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
#[cfg(not(target_os = "macos"))]
use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
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 crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png;
use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_png;
use image::imageops::FilterType;
use image::{DynamicImage, RgbImage};
use log::{error, info};
use ml_weights_registry::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
use std::path::PathBuf;
use tauri::{AppHandle, State};
const PROMPT_FILENAME : &str = "prompt.txt";
const PROMPT: &str = "shiba inu, cute dog, detailed, walking in a wooded forest, photorealistic, 8k";
@@ -42,28 +47,30 @@ pub async fn infer_image(
.map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
// TODO(bt,2025-02-17): Running out of vram with full image buffer size
let image = image.resize(512, 512, FilterType::CatmullRom);
let mut image = image.resize(512, 512, FilterType::CatmullRom);
let result = do_infer_image(
&prompt,
image,
strength,
&model_config,
&model_cache,
prompt_cache,
app,
&app_data_root
).await;
let image = match result {
Ok(image) => image,
Err(err) => {
error!("There was an error generating the image: {:?}", err);
return Err(format!("There was an error generating the image: {}", err));
}
};
#[cfg(not(target_os = "macos"))]
{
let result = infer_image_impl(
&prompt,
image,
strength,
&model_config,
&model_cache,
prompt_cache,
&app_data_root
).await;
let bytes = encode_rgb_image_base64_png(image)
image = match result {
Ok(img) => DynamicImage::ImageRgb8(img),
Err(err) => {
error!("There was an error generating the image: {:?}", err);
return Err(format!("There was an error generating the image: {}", err));
}
};
}
let bytes = encode_dynamic_image_base64_png(image)
.map_err(|err| format!("failure to encode image: {:?}", err))?;
info!("Inference successful; image converted to base64, serving back to browser...");
@@ -71,29 +78,40 @@ pub async fn infer_image(
Ok(bytes)
}
async fn do_infer_image(
#[cfg(not(target_os = "macos"))]
async fn infer_image_impl(
prompt: &str,
image: DynamicImage,
strength: Option<f64>,
config: &AppConfig,
model_cache: &ModelCache,
prompt_cache: State<'_, PromptCache>,
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(&LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16);
let clip_weights_path= weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER);
let args = Args {
image: &image,
prompt: prompt.to_string(),
uncond_prompt: NEGATIVE_PROMPT.to_string(),
model_cache,
configs: config,
prompt_cache: &prompt_cache,
i2i_strength: strength,
img2img_strength: strength,
cfg_scale: config.cfg_scale,
app: &app,
app_data_root,
use_flash_attn: 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,
};
let image = lcm_pipeline(args)
@@ -1,43 +1,31 @@
//! Adapted from https://github.com/dnanhkhoa/rust-background-removal
use crate::ml::background_removal::onnx_session::onnx_session;
use crate::ml::model_cache::ModelCache;
#[cfg(not(target_os = "macos"))]
use ml_models::ml::background_removal::remove_image_background::remove_image_background;
use crate::state::app_dir::AppDataRoot;
use crate::utils::image::decode_base64_image::decode_base64_image;
use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png;
use image::{imageops, DynamicImage};
use log::{error, info};
use ndarray::{Array, CowArray};
use ort::Value;
use tauri::{AppHandle, State};
use crate::ml::weights_registry::weights::{DIS_MEDIUM_ONNX, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
use log::info;
use ml_weights_registry::weights_registry::weights::DIS_MEDIUM_ONNX;
use tauri::State;
/// This handler removes the background from an image.
#[tauri::command]
pub async fn remove_background(
image: &str,
model_cache: State<'_, ModelCache>,
app_data_root: State<'_, AppDataRoot>,
app: AppHandle,
) -> Result<String, String> {
info!("remove_background endpoint called.");
let image = decode_base64_image(image)
.map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
let mut image = decode_base64_image(image)
.map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
// NB: Real time needed cuDNN 9, but this uses cuDNN 8:
// sudo apt-get install libcudnn8-dev libcudnn8
// update-alternatives: warning: forcing reinstallation of alternative /usr/include/x86_64-linux-gnu/cudnn_v9.h because link group libcudnn is broken
// update-alternatives: using /usr/include/x86_64-linux-gnu/cudnn_v8.h to provide /usr/include/cudnn.h (libcudnn) in manual mode
let model_path = app_data_root.weights_dir().weight_path(&DIS_MEDIUM_ONNX);
let session = onnx_session(model_path.to_str().unwrap()) // TODO
.map_err(|err| format!("failure to start onnx session: {:?}", err))?;
let image = process_dynamic_image(&session, image)
.map_err(|err| {
error!("Couldn't remove background from image: {}", err);
format!("Couldn't remove background from image: {}", err)
})?;
#[cfg(not(target_os = "macos"))]
{
image = remove_background_impl(&app_data_root, image).await?;
}
let image = encode_dynamic_image_base64_png(image)
.map_err(|err| format!("failure to encode image: {:?}", err))?;
@@ -45,48 +33,13 @@ pub async fn remove_background(
Ok(image)
}
fn process_dynamic_image(
session: &ort::Session,
dynamic_img: DynamicImage,
) -> Result<DynamicImage, anyhow::Error> {
let input_shape = session.inputs[0]
.dimensions()
.map(|dim| dim.unwrap())
.collect::<Vec<usize>>();
let input_img = dynamic_img.into_rgba8();
let scaling_factor = f32::min(
1., // Avoid upscaling
f32::min(
input_shape[3] as f32 / input_img.width() as f32, // Width ratio
input_shape[2] as f32 / input_img.height() as f32, // Height ratio
),
);
let mut resized_img = imageops::resize(
&input_img,
input_shape[3] as u32,
input_shape[2] as u32,
imageops::FilterType::Triangle,
);
let input_tensor = CowArray::from(
Array::from_shape_fn(input_shape, |indices| {
let mean = 128.;
let std = 256.;
#[cfg(not(target_os = "macos"))]
async fn remove_background_impl(app_data_root: &AppDataRoot, image: DynamicImage) -> Result<DynamicImage, String> {
let model_path = app_data_root.weights_dir().weight_path(&DIS_MEDIUM_ONNX);
(resized_img[(indices[3] as u32, indices[2] as u32)][indices[1]] as f32 - mean) / std
})
.into_dyn(),
);
let inputs = vec![Value::from_array(session.allocator(), &input_tensor)?];
let outputs = session.run(inputs)?;
let output_tensor = outputs[0].try_extract::<f32>()?;
for (indices, alpha) in output_tensor.view().indexed_iter() {
resized_img[(indices[3] as u32, indices[2] as u32)][3] = (alpha * 255.) as u8;
}
let output_img = imageops::resize(
&resized_img,
(input_img.width() as f32 * scaling_factor) as u32,
(input_img.height() as f32 * scaling_factor) as u32,
imageops::FilterType::Triangle,
);
Ok(DynamicImage::ImageRgba8(output_img))
}
let image = remove_image_background(model_path, image)
.await
.map_err(|err| format!("Couldn't remove image from image: {}", err))?;
Ok(image)
}
@@ -1,24 +0,0 @@
use std::sync::{Arc, RwLock};
use once_cell::sync::Lazy;
static COUNTER : Lazy<Arc<RwLock<u64>>> = Lazy::new(|| Arc::new(RwLock::new(0)));
#[tauri::command]
pub fn test_counter() -> u64 {
let value : u64;
{
match COUNTER.write() {
Ok(mut counter) => {
*counter += 1;
value = *counter;
},
Err(_e) => {
value = 0;
}
}
}
println!("I was invoked from JavaScript! {:?}", value);
value
}
@@ -1,21 +1,23 @@
use crate::ml::model_cache::ModelCache;
use crate::ml::prompt_cache::PromptCache;
use crate::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
#[cfg(not(target_os = "macos"))]
use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
use crate::state::app_config::AppConfig;
use crate::state::app_dir::AppDataRoot;
use crate::utils::image::decode_base64_image::decode_base64_image;
use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_png;
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 once_cell::sync::Lazy;
use std::io::Cursor;
use std::path::PathBuf;
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.
@@ -38,28 +40,31 @@ pub async fn text_to_image(
) -> Result<String, String> {
info!("text_to_image endpoint called.");
let image = IMAGE.resize(512, 512, FilterType::CatmullRom);
let mut image = IMAGE.resize(512, 512, FilterType::CatmullRom);
let result = do_infer_image(
&prompt,
image,
Some(STRENGTH),
&model_config,
&model_cache,
prompt_cache,
app,
&app_data_root
).await;
let image = match result {
Ok(image) => image,
Err(err) => {
error!("There was an error generating the image: {:?}", err);
return Err(format!("There was an error generating the image: {}", err));
}
};
#[cfg(not(target_os = "macos"))]
{
let result = text_to_image_impl(
&prompt,
image,
Some(STRENGTH),
&model_config,
&model_cache,
prompt_cache,
app,
&app_data_root
).await;
let bytes = encode_rgb_image_base64_png(image)
image = match result {
Ok(img) => DynamicImage::ImageRgb8(img),
Err(err) => {
error!("There was an error generating the image: {:?}", err);
return Err(format!("There was an error generating the image: {}", err));
}
};
}
let bytes = encode_dynamic_image_base64_png(image)
.map_err(|err| format!("failure to encode image: {:?}", err))?;
info!("Inference successful; image converted to base64, serving back to browser...");
@@ -67,7 +72,8 @@ pub async fn text_to_image(
Ok(bytes)
}
async fn do_infer_image(
#[cfg(not(target_os = "macos"))]
async fn text_to_image_impl(
prompt: &str,
image: DynamicImage,
strength: Option<f64>,
@@ -77,19 +83,29 @@ async fn do_infer_image(
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 {
image: &image,
prompt: prompt.to_string(),
uncond_prompt: "".to_string(),
model_cache,
configs: config,
prompt_cache: &prompt_cache,
i2i_strength: strength,
img2img_strength: strength,
cfg_scale: config.cfg_scale,
app: &app,
app_data_root,
use_flash_attn: 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,
};
let image = lcm_pipeline(args)
@@ -1,6 +1,5 @@
use serde::Serialize;
#[derive(Clone, Copy, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum NotificationModelType {
+9 -6
View File
@@ -1,21 +1,22 @@
pub mod endpoints;
pub mod events;
pub mod ml;
pub mod state;
pub mod stubs;
pub mod threads;
pub mod transfer;
pub mod utils;
use crate::endpoints::download_models::download_models;
use crate::endpoints::inpaint_image_endpoint::inpaint_image;
use crate::endpoints::flip_image::flip_image;
use crate::endpoints::realtime_image_endpoint::infer_image;
use crate::endpoints::remove_background_endpoint::remove_background;
use crate::endpoints::test_counter::test_counter;
use crate::endpoints::text_to_image_endpoint::text_to_image;
use crate::ml::model_cache::ModelCache;
use crate::ml::prompt_cache::PromptCache;
use crate::state::app_config::AppConfig;
use crate::stubs::model_cache::ModelCache;
use crate::stubs::prompt_cache::PromptCache;
use crate::threads::downloader_thread::downloader_thread;
use crate::utils::log_environment_details::log_environment_details;
use tauri_plugin_log::Target;
use tauri_plugin_log::TargetKind;
@@ -28,6 +29,7 @@ pub fn run() {
let config = AppConfig::init()
.expect("config should load");
let prompt_cache = PromptCache::with_capacity(8)
.expect("prompt cache should load");
@@ -68,9 +70,10 @@ pub fn run() {
.manage(app_data_root)
.invoke_handler(tauri::generate_handler![
download_models,
flip_image,
infer_image,
inpaint_image,
remove_background,
test_counter,
text_to_image,
])
.run(tauri::generate_context!())
@@ -1 +0,0 @@
pub mod onnx_session;
@@ -1,32 +1,14 @@
use crate::ml::model_file::StableDiffusionVersion;
use crate::state::app_dir::AppDataRoot;
use crate::state::os_platform::OsPlatform;
use crate::state::yaml_config::YamlConfig;
use candle_core::{DType, Device};
use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
use filesys::recursively_find_file_by_name::recursively_find_file_by_name;
use hf_hub::api::sync::Api;
use log::info;
use std::path::PathBuf;
use crate::stubs::model_config::ModelConfig;
const DEFAULT_SD_IMAGE_WIDTH: usize = 1024;
const DEFAULT_SD_IMAGE_HEIGHT: usize = 1024;
const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32;
const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::V1_5;
/// A central place to configure app-wide details.
pub struct AppConfig {
/// We only support one device type currently, and we can't multiplex
/// over several devices yet.
pub device: Device,
/// The core numeric type to use for models.
pub dtype: DType,
pub sd_version: StableDiffusionVersion,
pub sd_config: StableDiffusionConfig,
pub image_height: usize,
pub image_width: usize,
pub scheduler_steps: usize,
@@ -36,11 +18,10 @@ pub struct AppConfig {
pub cfg_scale: Option<f64>,
/// Probably shouldn't be here.
pub hf_api: Api,
/// Root location for application data.
pub app_data_root: AppDataRoot,
pub model_config: ModelConfig,
}
impl AppConfig {
@@ -50,30 +31,6 @@ impl AppConfig {
println!("Configs: {:?}", yaml_configs);
let device = Device::new_cuda(0)
.unwrap_or_else(|e| {
info!("CUDA not available ({}), falling back to CPU", e);
Device::Cpu
});
let sd_config = match SD_VERSION {
StableDiffusionVersion::Turbo => {
StableDiffusionConfig::sdxl_turbo(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
}
StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => {
StableDiffusionConfig::v1_5(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
}
StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint => {
StableDiffusionConfig::v2_1(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
}
StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint => {
StableDiffusionConfig::sdxl(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
}
};
let hf_api = Api::new()?;
let os_platform = OsPlatform::get();
let default_data_root = match os_platform {
@@ -92,18 +49,14 @@ impl AppConfig {
println!("App data root: {:?}", app_data_root.path());
Ok(Self {
device,
dtype: DEFAULT_SD_NUMERIC_DATATYPE,
sd_version: SD_VERSION,
sd_config,
image_height: yaml_configs.image_height.unwrap_or(DEFAULT_SD_IMAGE_HEIGHT),
image_width: yaml_configs.image_width.unwrap_or(DEFAULT_SD_IMAGE_WIDTH),
scheduler_steps: yaml_configs.scheduler_steps.unwrap_or(4),
scheduler_samples: yaml_configs.scheduler_samples.unwrap_or(15),
seed: yaml_configs.seed,
cfg_scale: yaml_configs.cfg_scale,
hf_api,
app_data_root,
model_config: ModelConfig::init()?,
})
}
}
@@ -1,8 +1,8 @@
use crate::ml::weights_registry::weight_descriptor::WeightDescriptor;
use crate::state::expanduser::expanduser;
use crate::state::os_platform::OsPlatform;
use anyhow::anyhow;
use directories::UserDirs;
use ml_weights_registry::weights_registry::weight_descriptor::WeightDescriptor;
use std::path::{Path, PathBuf};
use tempdir::TempDir;
use tempfile::{Builder, NamedTempFile};
@@ -0,0 +1,3 @@
pub mod model_cache;
pub mod model_config;
pub mod prompt_cache;
@@ -0,0 +1,14 @@
#[cfg(not(target_os = "macos"))]
pub use ml_models::ml::model_cache::ModelCache;
#[cfg(not(not(target_os = "macos")))]
pub struct ModelCache {
// Intentionally left blank
}
#[cfg(not(not(target_os = "macos")))]
impl ModelCache {
pub fn new() -> Self {
Self {}
}
}
@@ -0,0 +1,14 @@
#[cfg(not(target_os = "macos"))]
pub use ml_models::ml::model_config::ModelConfig;
#[cfg(not(not(target_os = "macos")))]
pub struct ModelConfig {
// Intentionally left blank
}
#[cfg(not(not(target_os = "macos")))]
impl ModelConfig {
pub fn init() -> anyhow::Result<Self> {
Ok(Self {})
}
}
@@ -0,0 +1,15 @@
#[cfg(not(target_os = "macos"))]
pub use ml_models::ml::prompt_cache::PromptCache;
#[cfg(not(not(target_os = "macos")))]
pub struct PromptCache {
// Intentionally left blank
}
#[cfg(not(not(target_os = "macos")))]
impl PromptCache {
pub fn with_capacity(capacity: usize) -> anyhow::Result<Self> {
Ok(Self {})
}
}
@@ -1,9 +1,11 @@
use crate::events::notification_event::{NotificationEvent, NotificationModelType};
use crate::state::app_dir::AppDataRoot;
use crate::transfer::download::{download_async, ProgressUpdate};
use futures::stream::FuturesUnordered;
use futures::StreamExt;
use futures::stream::FuturesUnordered;
use log::{debug, error, info};
use ml_weights_registry::weights_registry::weight_descriptor::WeightFunction;
use ml_weights_registry::weights_registry::weights::{CLIP_JSON, DIS_MEDIUM_ONNX, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2, SDXL_TURBO_UNET, SDXL_TURBO_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
use std::collections::VecDeque;
use std::path::PathBuf;
use std::sync::mpsc::Receiver;
@@ -11,8 +13,6 @@ use std::sync::{Arc, Mutex};
use tauri::{AppHandle, Emitter};
use tempfile::NamedTempFile;
use tokio::task::JoinHandle;
use crate::ml::weights_registry::weight_descriptor::WeightFunction;
use crate::ml::weights_registry::weights::{CLIP_JSON, DIS_MEDIUM_ONNX, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2, SDXL_TURBO_UNET, SDXL_TURBO_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
const NOTIFICATION_CHANNEL_NAME : &str = "notification";
const MAX_FILES : usize = 8;
@@ -1,7 +1,6 @@
use crate::ml::image::tensor_to_image_buffer::RgbImage;
use base64::prelude::BASE64_STANDARD;
use base64::Engine;
use image::ImageFormat;
use image::{ImageFormat, RgbImage};
use std::io::Cursor;
pub fn encode_rgb_image_base64_png(
@@ -1,11 +1,11 @@
use log::info;
use nvidia_checker::{get_cuda_version, get_cudnn_version, get_kernel_version, get_nvidia_driver_version, get_os_version, get_tensorrt_version};
//use nvidia_checker::{get_cuda_version, get_cudnn_version, get_kernel_version, get_nvidia_driver_version, get_os_version, get_tensorrt_version};
pub fn log_environment_details() {
info!("NVIDIA driver version: {}", get_nvidia_driver_version(false));
info!("CUDA version: {}", get_cuda_version(false));
info!("cuDNN version: {}", get_cudnn_version(false));
info!("TensorRT version: {}", get_tensorrt_version(false));
info!("OS version: {}", get_os_version(false));
info!("Kernel version: {}", get_kernel_version(false));
//info!("NVIDIA driver version: {}", get_nvidia_driver_version(false));
//info!("CUDA version: {}", get_cuda_version(false));
//info!("cuDNN version: {}", get_cudnn_version(false));
//info!("TensorRT version: {}", get_tensorrt_version(false));
//info!("OS version: {}", get_os_version(false));
//info!("Kernel version: {}", get_kernel_version(false));
}
+1 -4
View File
@@ -32,9 +32,6 @@
"icons/128x128@2x.png",
"icons/icon.icns",
"icons/icon.ico"
],
"resources": {
"../../../windows_dlls/*.dll": "."
}
]
}
}
+18 -6
View File
@@ -8,6 +8,7 @@
"name": "storyteller-board",
"version": "0.0.0",
"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",
@@ -381,6 +382,21 @@
"node": ">=6.9.0"
}
},
"node_modules/@devnomic/marquee": {
"version": "1.0.3",
"resolved": "https://registry.npmjs.org/@devnomic/marquee/-/marquee-1.0.3.tgz",
"integrity": "sha512-q8QUtNO8ff+UvfrKJsHBtxh4fbOfhGk6SSdK0Uu/XZFdEzrmiGJ/l0FuVHRYsPAmzgnhVVnz1dGNB5CWGLMVvw==",
"dependencies": {
"clsx": "^2.1.0",
"tailwind-merge": "^2.2.0"
},
"peerDependencies": {
"@types/react": "^18.2.47",
"@types/react-dom": "^18.2.18",
"react": "^18.2.0",
"react-dom": "^18.2.0"
}
},
"node_modules/@emotion/is-prop-valid": {
"version": "1.3.0",
"resolved": "https://registry.npmjs.org/@emotion/is-prop-valid/-/is-prop-valid-1.3.0.tgz",
@@ -1510,14 +1526,12 @@
"node_modules/@types/prop-types": {
"version": "15.7.12",
"resolved": "https://registry.npmjs.org/@types/prop-types/-/prop-types-15.7.12.tgz",
"integrity": "sha512-5zvhXYtRNRluoE/jAp4GVsSduVUzNWKkOZrCDBWYtE7biZywwdC2AcEzg+cSMLFRfVgeAFqpfNabiPjxFddV1Q==",
"dev": true
"integrity": "sha512-5zvhXYtRNRluoE/jAp4GVsSduVUzNWKkOZrCDBWYtE7biZywwdC2AcEzg+cSMLFRfVgeAFqpfNabiPjxFddV1Q=="
},
"node_modules/@types/react": {
"version": "18.3.4",
"resolved": "https://registry.npmjs.org/@types/react/-/react-18.3.4.tgz",
"integrity": "sha512-J7W30FTdfCxDDjmfRM+/JqLHBIyl7xUIp9kwK637FGmY7+mkSFSe6L4jpZzhj5QMfLssSDP4/i75AKkrdC7/Jw==",
"dev": true,
"dependencies": {
"@types/prop-types": "*",
"csstype": "^3.0.2"
@@ -1527,7 +1541,6 @@
"version": "18.3.0",
"resolved": "https://registry.npmjs.org/@types/react-dom/-/react-dom-18.3.0.tgz",
"integrity": "sha512-EhwApuTmMBmXuFOikhQLIBUn6uFg81SwLMOAUgodJF14SOBOCMdU04gDoYi0WOJJHD144TL32z4yDqCW3dnkQg==",
"dev": true,
"dependencies": {
"@types/react": "*"
}
@@ -2364,8 +2377,7 @@
"node_modules/csstype": {
"version": "3.1.3",
"resolved": "https://registry.npmjs.org/csstype/-/csstype-3.1.3.tgz",
"integrity": "sha512-M1uQkMl8rQK/szD0LNhtqxIPLpimGm8sOBwU7lLnCpSbTyY3yeU1Vc7l4KT5zT4s/yOxHH5O7tIuuLOCnLADRw==",
"dev": true
"integrity": "sha512-M1uQkMl8rQK/szD0LNhtqxIPLpimGm8sOBwU7lLnCpSbTyY3yeU1Vc7l4KT5zT4s/yOxHH5O7tIuuLOCnLADRw=="
},
"node_modules/data-view-buffer": {
"version": "1.0.1",
+1
View File
@@ -11,6 +11,7 @@
"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",
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.2 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 276 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 320 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 564 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 366 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 452 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 337 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 378 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 302 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 446 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 238 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 344 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.2 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.1 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 299 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 396 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 318 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 521 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 318 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 396 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 316 KiB

@@ -0,0 +1,25 @@
import { useCallback, useRef } from "react";
import { useRenderCounter } from "~/hooks/useRenderCounter";
import { EditEngine } from "~/KonvaApp/EditEngine";
import { KonvaCanvasContainer } from "~/KonvaRootComponent/KonvaCanvasContainer";
export const EditModeRootComponent = ({ className }: { className: string }) => {
useRenderCounter("EditModeRootComponent");
const editEngineRef = useRef<EditEngine | null>(null);
const konvaContainerCallbackRef = useCallback((node: HTMLDivElement) => {
if (node !== null && editEngineRef.current === null) {
editEngineRef.current = new EditEngine(node);
}
}, []);
return (
<>
<KonvaCanvasContainer
className={className}
ref={konvaContainerCallbackRef}
/>
</>
);
};
@@ -0,0 +1,145 @@
import { useEffect, useState } from "react";
import { Marquee } from "@devnomic/marquee";
import "@devnomic/marquee/dist/index.css";
function getRandomImages(count: number) {
const images = Array.from({ length: count }, (_, i) => ({
id: i + 1,
src: `/images/generate/image${i + 1}.png`,
alt: `Generated image ${i + 1}`,
}));
// Shuffle the array
return [...images].sort(() => Math.random() - 0.5);
}
interface ImageColumnProps {
images: Array<{ id: number; src: string; alt: string }>;
delay: number;
reverse?: boolean;
isLoaded: boolean;
duration?: number;
}
function ImageColumn({
images,
delay,
reverse = false,
duration = 100,
isLoaded,
}: ImageColumnProps) {
return (
<div className="flex h-screen flex-col gap-6">
<Marquee
fade={false}
direction="up"
reverse={reverse}
pauseOnHover={false}
className={`h-full gap-[24px]`}
style={{ "--duration": `${duration}s` } as React.CSSProperties}
innerClassName="gap-[24px] [--gap:24px]"
numberOfCopies={2}
>
{images.map((image, i) => (
<div
key={`${image.id}-${i}`}
className={`relative aspect-[4/5] w-full overflow-hidden rounded-lg shadow-lg transition-all duration-700 ease-out ${
isLoaded
? "translate-y-0 opacity-100"
: "translate-y-12 opacity-0"
}`}
style={{
transitionDelay: `${delay + i * 50}ms`,
}}
>
<img
src={image.src}
alt={image.alt}
className="w-full object-cover"
/>
</div>
))}
</Marquee>
</div>
);
}
export default function GalleryBackground() {
const [images, _setImages] = useState(() => getRandomImages(20));
const [isLoaded, setIsLoaded] = useState(false);
useEffect(() => {
// delay for smooth animation
const timer = setTimeout(() => {
setIsLoaded(true);
}, 100);
return () => clearTimeout(timer);
}, []);
// Columns for masonry layout
const columns = {
col1: images.filter((_, i) => i % 5 === 0),
col2: images.filter((_, i) => i % 5 === 1),
col3: images.filter((_, i) => i % 5 === 2),
col4: images.filter((_, i) => i % 5 === 3),
col5: images.filter((_, i) => i % 5 === 4),
};
return (
<>
<div className="fixed inset-0 z-[1] overflow-hidden bg-[radial-gradient(50%_50%_at_50%_50%,rgba(0,0,0,0.00)_49%,rgba(0,0,0,0.60)_100%)]" />
<div className="fixed inset-0 z-0 overflow-hidden">
<div
className={`h-full w-full transition-opacity duration-1000 ${isLoaded ? "opacity-50" : "opacity-0"}`}
>
<div className="absolute inset-0 grid grid-cols-5 gap-6 px-6 opacity-35">
<ImageColumn
images={columns.col1}
delay={0}
isLoaded={isLoaded}
duration={100}
/>
<ImageColumn
images={columns.col2}
delay={columns.col1.length * 50}
isLoaded={isLoaded}
reverse={true}
duration={100}
/>
<ImageColumn
images={columns.col3}
delay={(columns.col1.length + columns.col2.length) * 50}
isLoaded={isLoaded}
duration={125}
/>
<ImageColumn
images={columns.col4}
delay={
(columns.col1.length +
columns.col2.length +
columns.col3.length) *
50
}
isLoaded={isLoaded}
reverse={true}
duration={160}
/>
<ImageColumn
images={columns.col5}
delay={
(columns.col1.length +
columns.col2.length +
columns.col3.length +
columns.col4.length) *
50
}
isLoaded={isLoaded}
duration={110}
/>
</div>
</div>
</div>
</>
);
}
@@ -0,0 +1,196 @@
import { faSparkles } from "@fortawesome/pro-solid-svg-icons";
import { useSignals } from "@preact/signals-react/runtime";
import { Button, LoadingSpinner } from "~/components/ui";
import { useRenderCounter } from "~/hooks/useRenderCounter";
import { GenerationLoadingState, generationSignal } from "~/signals";
import { useRef, useState } from "react";
import BackgroundGallery from "./BackgroundGallery";
import { twMerge } from "tailwind-merge";
import { Transition } from "@headlessui/react";
import { GenerationEngine } from "~/KonvaApp/GenerationEngine";
import { ensureBase64Prefix } from "~/KonvaApp/EngineUtitlities/Base64Helpers";
export const GenerationRootComponent = ({
generationEngineRef,
}: {
generationEngineRef: React.MutableRefObject<GenerationEngine | null>;
}) => {
// This is a hook that will log the number of times the component has rerendered
// Let's make sure we only log once
useRenderCounter("GenerationRootComponent");
useSignals();
const generationState = generationSignal.value;
const inputRef = useRef<HTMLTextAreaElement>(null);
const [isInputFocused, setIsInputFocused] = useState(false);
if (generationEngineRef.current === null) {
generationEngineRef.current = new GenerationEngine();
}
const handleGenerate = () => {
if (!inputRef.current) {
return;
}
const prompt = inputRef.current.value;
generationSignal.value = {
loadingState: GenerationLoadingState.GENERATING,
prompt,
};
};
let contentElement;
switch (generationState.loadingState) {
case GenerationLoadingState.INIT:
contentElement = null;
break;
case GenerationLoadingState.GENERATING:
case GenerationLoadingState.GENERATED:
contentElement = (
<GenerationContent
loadingState={generationState.loadingState}
imageB64={generationState.imageB64}
prompt={generationState.prompt}
/>
);
break;
}
// const isButtonDisabled = generationState.loadingState === GenerationLoadingState.GENERATING;
const isButtonDisabled = false;
return (
<>
<div className="fixed z-10 flex h-full w-full flex-col items-center justify-center transition-all duration-500">
<Transition
show={generationState.loadingState !== GenerationLoadingState.INIT}
enter="transition-opacity duration-500"
enterFrom="opacity-0"
enterTo="opacity-100"
leave="transition-opacity duration-500"
leaveFrom="opacity-100"
leaveTo="opacity-0"
as="div"
>
{contentElement}
</Transition>
<div
className={twMerge(
"absolute flex w-full flex-col items-center justify-center gap-16 transition-all duration-500",
generationState.loadingState === GenerationLoadingState.INIT
? "bottom-1/2 translate-y-1/2 transform"
: "bottom-8",
)}
>
<Transition
show={generationState.loadingState === GenerationLoadingState.INIT}
enter="transition-opacity duration-200"
enterFrom="opacity-0"
enterTo="opacity-100"
leave="transition-opacity duration-200"
leaveFrom="opacity-100"
leaveTo="opacity-0"
as="div"
>
<div className="flex flex-col items-center justify-center text-center drop-shadow-xl">
<span className="text-8xl font-bold">Generate Image</span>
<span className="pt-2 text-2xl opacity-80">
Imagine, Describe, Generate.
</span>
</div>
</Transition>
<div
className={twMerge(
"glass flex min-h-[56px] w-full max-w-[900px] items-end rounded-xl border-2 p-3 shadow-xl transition-all duration-[400ms] ease-in-out",
isInputFocused ? "border-primary-400/60" : "border-ui-panel",
)}
>
<textarea
ref={inputRef}
rows={1}
placeholder="Describe what you want to see..."
className="max-h-[120px] min-h-[40px] flex-1 resize-none overflow-y-auto rounded-md bg-transparent px-3 py-2 text-lg leading-normal focus:outline-none"
style={{ lineHeight: "24px" }}
onInput={(e) => {
const target = e.target as HTMLTextAreaElement;
target.style.height = "auto";
target.style.height = `${target.scrollHeight}px`;
}}
autoFocus={true}
onFocus={() => setIsInputFocused(true)}
onBlur={() => setIsInputFocused(false)}
/>
<Button
icon={faSparkles}
variant="primary"
className="text-md ml-2 self-end"
onClick={handleGenerate}
disabled={isButtonDisabled}
loading={
generationState.loadingState ===
GenerationLoadingState.GENERATING
}
>
Generate
</Button>
</div>
</div>
</div>
<Transition
show={generationState.loadingState === GenerationLoadingState.INIT}
enter="transition-opacity duration-700"
enterFrom="opacity-0"
enterTo="opacity-100"
leave="transition-opacity duration-700"
leaveFrom="opacity-100"
leaveTo="opacity-0"
as="div"
>
<BackgroundGallery />
</Transition>
</>
);
};
const GenerationContent = (generationState: {
loadingState: GenerationLoadingState;
imageB64?: string;
prompt: string;
}) => {
let imgBoxContent;
if (generationState.loadingState === GenerationLoadingState.GENERATING) {
imgBoxContent = (
<div className="flex h-full w-full flex-col items-center justify-center gap-7">
<LoadingSpinner isShowing={true} message="Generating image..." />
</div>
);
} else if (
generationState.loadingState === GenerationLoadingState.GENERATED &&
generationState.imageB64
) {
imgBoxContent = (
<img
src={ensureBase64Prefix(generationState.imageB64)}
alt="Generated Image"
className="h-full w-full object-contain"
/>
);
}
return (
<>
<div className="absolute inset-0 flex h-full w-full flex-col items-center justify-center">
<div className="aspect-[1/1] w-[min(80vw,68vh)] overflow-clip rounded-t-lg bg-[#29292D]/60">
{imgBoxContent}
</div>
<div className="mb-16 flex w-[min(80vw,68vh)] rounded-b-lg bg-[#29292D]">
<span className="p-3.5 text-sm opacity-50">
{generationState.prompt || "No prompt provided"}
</span>
</div>
</div>
</>
);
};
@@ -0,0 +1,365 @@
import Konva from "konva";
import { uiAccess } from "~/signals";
import { CommandManager, MatteBox } from "./EngineUtitlities";
import { ImageNode } from "./Nodes";
import { PaintNode } from "./Nodes/PaintNode";
import {
NodesManager,
NodesTranslationEventDetails,
NodeTransformationEventDetails,
NodeTransformer,
SelectionManager,
SelectionManagerEvents,
SelectorSquare,
} from "./NodesManagers";
import { UndoStackManager } from "./UndoRedo";
import { ToolbarMainButtonNames } from "~/components/features/ToolbarMain/enum";
import { EditModes, VideoResolutions } from "./constants";
import { ToolbarNodeButtonNames } from "~/components/features/ToolbarNode/enums";
import { RealTimeDrawEngine } from "./RenderingPrimitives/RealTimeDrawEngine";
export interface RenderingOptions {
artstyle: string;
positivePrompt: string;
negativePrompt: string;
cinematic: boolean;
enginePreProcessing: boolean;
faceDetail: boolean;
lipSync: boolean;
upscale: boolean;
styleStrength: number;
}
export class EditEngine {
private editMode: EditModes = EditModes.SELECT;
private boardCanvasRef: HTMLDivElement;
private stage: Konva.Stage;
private bgLayer: Konva.Layer;
private mediaLayer: Konva.Layer;
private nodeIsolationLayer: Konva.Layer;
private uiLayer: Konva.Layer;
private offScreenCanvas: OffscreenCanvas;
private realTimeDrawEngine: RealTimeDrawEngine;
private nodesManager: NodesManager;
private nodeTransformer: NodeTransformer;
private selectionManager: SelectionManager;
private selectorSquare: SelectorSquare;
private matteBox: MatteBox;
private undoStackManager: UndoStackManager;
private commandManager: CommandManager;
public segmentationButtonCanBePressed: boolean = true;
// signal reference
constructor(boardCanvasRef: HTMLDivElement) {
if (import.meta.env.DEV) {
console.log("Engine Created");
}
this.editMode = EditModes.SELECT;
this.boardCanvasRef = boardCanvasRef;
this.stage = new Konva.Stage({
container: this.boardCanvasRef,
width: window.innerWidth,
height: window.innerHeight,
});
this.bgLayer = new Konva.Layer();
this.mediaLayer = new Konva.Layer();
this.nodeIsolationLayer = new Konva.Layer();
this.uiLayer = new Konva.Layer();
this.stage.add(this.bgLayer);
this.stage.add(this.mediaLayer);
this.stage.add(this.nodeIsolationLayer);
this.stage.add(this.uiLayer);
// Konva Transformer
this.nodeTransformer = new NodeTransformer();
this.uiLayer.add(this.nodeTransformer.getKonvaNode());
// core layer for all the work done.
this.offScreenCanvas = new OffscreenCanvas(0, 0);
// Collection of all Nodes
this.nodesManager = new NodesManager();
// Partial Collection of selected Nodes
this.selectionManager = new SelectionManager({
nodeTransformerRef: this.nodeTransformer,
mediaLayerRef: this.mediaLayer,
});
this.realTimeDrawEngine = new RealTimeDrawEngine({
width: VideoResolutions.SQUARE_1024.width,
height: VideoResolutions.SQUARE_1024.height,
mediaLayerRef: this.mediaLayer,
offScreenCanvas: this.offScreenCanvas,
onDraw: async (canvas, lineBounds) => {
await this.addPaintNode(canvas, lineBounds);
},
});
// Selector Square to select Nodes
this.selectorSquare = new SelectorSquare({
captureCanvasRef: this.realTimeDrawEngine.captureCanvas,
mediaLayerRef: this.mediaLayer,
nodesManagerRef: this.nodesManager,
selectionManagerRef: this.selectionManager,
stageRef: this.stage,
});
this.uiLayer.add(this.selectorSquare.getKonvaNode());
//Collection of commands for undo-redo
this.undoStackManager = new UndoStackManager(() => {
this.realTimeDrawEngine.render();
});
this.commandManager = new CommandManager({
mediaLayerRef: this.mediaLayer,
nodesManagerRef: this.nodesManager,
nodeTransformerRef: this.nodeTransformer,
selectionManagerRef: this.selectionManager,
renderEngineRef: this.realTimeDrawEngine,
undoStackManagerRef: this.undoStackManager,
});
this.matteBox = new MatteBox({
boardCanvasSize: {
width: this.boardCanvasRef.clientWidth,
height: this.boardCanvasRef.clientHeight,
},
captureCanvasSize: this.realTimeDrawEngine.captureCanvas.getSize(),
uiLayerRef: this.uiLayer,
});
// some of the managers has events
// hence, lastly, setup these events
this.setupEventSystem();
this.setEditMode(this.editMode);
}
private setEditMode(newEditMode: EditModes) {
this.editMode = newEditMode;
switch (this.editMode) {
case EditModes.SELECT:
default: {
console.log("EDITMODE: SELECT");
this.realTimeDrawEngine.disablePaintMode();
this.selectorSquare.enable();
this.selectionManager.enable();
uiAccess.toolbarMain.enable();
uiAccess.toolbarMain.changeButtonState(ToolbarMainButtonNames.SELECT, {
active: true,
});
uiAccess.toolbarMain.changeButtonState(ToolbarMainButtonNames.PAINT, {
active: false,
});
this.matteBox.disable();
document.body.style.cursor = "default";
return;
}
case EditModes.EDIT: {
console.log("EDITMODE: EDIT");
this.realTimeDrawEngine.enablePaintMode();
this.selectorSquare.disable();
this.selectionManager.disable();
}
}
}
private setupEventSystem() {
// Listen to changes in container size
const resizeObserver = new ResizeObserver(() => {
this.onBoardCanvasResize();
});
resizeObserver.observe(this.boardCanvasRef);
this.onBoardCanvasResize();
// Listen to Nodes
this.selectionManager.eventTarget.addEventListener(
SelectionManagerEvents.NODES_TRANSLATIONS,
((event: CustomEvent<NodesTranslationEventDetails>) => {
//console.log("Event: SelectionManager -> Engine", event);
this.commandManager.translateNodes(event.detail);
}) as EventListener,
);
this.selectionManager.eventTarget.addEventListener(
SelectionManagerEvents.NODES_TRANSFORMATION,
((event: CustomEvent<NodeTransformationEventDetails>) => {
//console.log("Event: SelectionManager -> Engine", event);
this.commandManager.transformNodes(event.detail);
}) as EventListener,
);
// Listen to Toolbar Main
// uiEvents.toolbarMain.UNDO.onClick(() => {
// this.undoStackManager.undo();
// });
// uiEvents.toolbarMain.REDO.onClick(() => {
// this.undoStackManager.redo();
// this.realTimeDrawEngine.render();
// });
// uiEvents.toolbarMain.SAVE.onClick(async (/*event*/) => {
// await this.realTimeDrawEngine.saveOutput();
// });
// uiEvents.toolbarMain.SELECT.onClick(() => {
// console.log("Toolbar Main >> Select");
// this.setAppMode(AppModes.SELECT);
// });
// uiEvents.toolbarMain.PAINT.onClick(() => {
// console.log("Toolbar Main >> Paint");
// this.setAppMode(AppModes.PAINT);
// });
}
disableAllButtons() {
const buttonNames = Object.values(ToolbarNodeButtonNames);
for (const name of buttonNames) {
uiAccess.toolbarNode.changeButtonState(name, { disabled: true });
}
}
async enableAllButtons() {
const buttonNames = Object.values(ToolbarNodeButtonNames);
for (const name of buttonNames) {
await uiAccess.toolbarNode.changeButtonState(name, { disabled: false });
}
}
sleep(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
private onBoardCanvasResize() {
this.realTimeDrawEngine.updateCaptureCanvas(undefined, undefined);
this.matteBox.updateSize({
boardCanvasSize: {
width: this.boardCanvasRef.offsetWidth,
height: this.boardCanvasRef.offsetHeight,
},
});
this.uiLayer.draw();
this.stage.width(this.boardCanvasRef.offsetWidth);
this.stage.height(this.boardCanvasRef.offsetHeight);
this.stage.draw(); // Redraw the canvas
}
// Sandbox is quickly a way to test your idea.
public async sandbox() {}
public onMessage(event: MessageEvent) {
console.log("Message From Shared Worker");
console.log(event);
}
public initializeStage() {
// load canvas that was originaly saved TODO Save manager for resharing.
uiAccess.toolbarNode.hide();
uiAccess.loadingBar.hide();
this.setupStage();
}
public isInitialized() {
return this.stage !== null;
}
public async setupStage() {
// Frame rate inicator
const textNode = new Konva.Text({
x: 10,
y: 80,
text: "",
fontSize: 18,
fontFamily: "Source Sans 3",
fill: "white",
});
const anim = new Konva.Animation((frame) => {
if (frame && import.meta.env.DEV) {
const timeDiff = frame.timeDiff;
const frameRate = frame.frameRate;
textNode.setText(
`FrameTime: ${timeDiff.toFixed(0)} ms\nFrameRate: ${frameRate.toFixed(0)} fps`,
);
}
}, this.mediaLayer);
anim.start();
this.uiLayer.add(textNode);
this.addKeyboardShortcuts();
}
public addImage(imageFile: File) {
const imageNode = new ImageNode({
mediaLayerRef: this.mediaLayer,
canvasPosition: this.realTimeDrawEngine.captureCanvas.position(),
canvasSize: this.realTimeDrawEngine.captureCanvas.size(),
imageFile: imageFile,
selectionManagerRef: this.selectionManager,
loaded: async () => {
await this.realTimeDrawEngine.render();
this.setEditMode(EditModes.SELECT);
},
});
this.commandManager.createNode(imageNode);
}
public addImageFromImageBitmap(imageBitmap: ImageBitmap) {
const imageNode = new ImageNode({
mediaLayerRef: this.mediaLayer,
canvasPosition: this.realTimeDrawEngine.captureCanvas.position(),
canvasSize: this.realTimeDrawEngine.captureCanvas.size(),
imageBitmap: imageBitmap,
selectionManagerRef: this.selectionManager,
loaded: async () => {
await this.realTimeDrawEngine.render();
this.setEditMode(EditModes.SELECT);
},
});
this.commandManager.createNode(imageNode);
}
public addPaintNode(
canvas: HTMLCanvasElement,
lineBounds: {
width: number;
height: number;
x: number;
y: number;
},
) {
var node = new PaintNode({
canvasElement: canvas,
lineBounds: lineBounds,
mediaLayerRef: this.mediaLayer,
selectionManagerRef: this.selectionManager,
loaded: async () => {
await this.realTimeDrawEngine.render();
},
});
this.commandManager.createNode(node);
}
// Events for Undo and Redo
private addKeyboardShortcuts() {
window.addEventListener("keydown", (event) => {
if (event.ctrlKey && event.key === "z") {
this.undoStackManager.undo();
} else if (
(event.ctrlKey && event.key === "y") ||
(event.ctrlKey && event.shiftKey && event.key === "Z")
) {
this.undoStackManager.redo();
} else if (event.key === "Delete") {
this.commandManager.deleteNodes();
}
});
}
}
@@ -0,0 +1,7 @@
const ensureBase64Prefix = (base64String: string) => {
return base64String.startsWith("data:")
? base64String
: `data:image/png;base64,${base64String}`;
}
export { ensureBase64Prefix };
@@ -0,0 +1,46 @@
import { invoke } from "@tauri-apps/api/core";
import {
cancelGeneration,
finishGeneration,
onStartGeneration,
} from "~/signals/generation/generationSignals";
export class GenerationEngine {
private isGenerating = false;
constructor() {
console.log("GenerationEngine constructor");
this.attachEventListeners();
}
private attachEventListeners() {
onStartGeneration((prompt) => {
this.invokeGeneration(prompt);
});
}
private async invokeGeneration(prompt: string) {
if (this.isGenerating) {
return;
}
this.isGenerating = true;
console.log("Invoking generation for prompt:", prompt);
try {
const base64BitmapResponse = await invoke("text_to_image", {
prompt: prompt,
});
console.log("Generation finished with response:", base64BitmapResponse);
const base64string = base64BitmapResponse as string;
finishGeneration(base64string);
} catch (error) {
console.error("Error during image processing:", error);
cancelGeneration();
} finally {
this.isGenerating = false;
}
}
}
@@ -25,6 +25,7 @@ import {
isLoadingVisible,
loadingProgress,
} from "~/signals/uiEvents/loadingIndicator";
import { ensureBase64Prefix } from "../EngineUtitlities/Base64Helpers";
import {DecodeBase64ToImage} from "~/utilities/DecodeBase64ToImage.ts";
interface ServerSettings {
@@ -7,6 +7,11 @@ export enum AppModes {
RENDERING = "RENDERING",
}
export enum EditModes {
SELECT = "SELECT",
EDIT = "EDIT",
}
export const VideoResolutions = {
VERTICAL_720: { width: 720, height: 1280 },
LANDSCAPE_720: { width: 1280, height: 720 },
+7
View File
@@ -1,4 +1,5 @@
// to do fix https://codesandbox.io/p/sandbox/react-konva-infinite-grid-kkndq?file=%2Fsrc%2Findex.js the dotted background doesn't move when draggable.
import { EditEngine } from "./EditEngine";
import { Engine } from "./Engine";
import { EngineOptions } from "./types";
export { Engine as EngineType };
@@ -8,3 +9,9 @@ export const KonvaApp = (element: HTMLDivElement, options: EngineOptions) => {
engine.initializeStage(options.sceneToken);
return engine;
};
export const KonvaEditModeApp = (element: HTMLDivElement) => {
const engine = new EditEngine(element);
engine.initializeStage();
return engine;
};
+40 -4
View File
@@ -10,6 +10,13 @@ import { KonvaRootComponent } from "~/KonvaRootComponent";
import { ToolbarTopLeft } from "~/components/features/ToolbarTopLeft";
import { ToolbarTopRight } from "~/components/features/ToolbarTopRight";
import { ToolbarTopCenter } from "~/components/features/ToolbarTopCenter";
import { appMode } from "~/signals";
import { useSignals } from "@preact/signals-react/runtime";
import { GenerationRootComponent } from "~/GenerationRootComponent/GenerationRootComponent";
import { EditModeRootComponent } from "~/EditModeRootComponent/EditModeRootComponent";
import { useRef } from "react";
import { EngineType } from "~/KonvaApp";
import { GenerationEngine } from "~/KonvaApp/GenerationEngine";
export const Main = withProtectionRoute(() => {
// This is a hook that will log the number of times the component has rerendered
@@ -19,12 +26,41 @@ export const Main = withProtectionRoute(() => {
if (import.meta.env.DEV && sceneToken === "debug") {
sceneToken = "m_p8nkry6m5j22w586xyex0w4a4pznbx";
}
useSignals();
const realtimeEngineRef = useRef<EngineType | null>(null);
const generationEngineRef = useRef<GenerationEngine | null>(null);
const appModeValue = appMode.value;
let childView;
switch (appModeValue) {
case "edit":
childView = (
<EditModeRootComponent className="col-span-12 col-start-1 row-span-12 row-start-1" />
);
break;
case "generate":
childView = (
<GenerationRootComponent generationEngineRef={generationEngineRef} />
);
break;
case "realtime":
default:
childView = (
<KonvaRootComponent
sceneToken={sceneToken}
className="col-span-12 col-start-1 row-span-12 row-start-1"
/>
);
break;
}
return (
<div className="fixed grid h-full w-full grid-cols-12 grid-rows-12">
<KonvaRootComponent
sceneToken={sceneToken}
className="col-span-12 col-start-1 row-span-12 row-start-1"
/>
{childView}
<div className="absolute top-0 w-full p-3.5">
<div className="relative flex w-full">
<div className="absolute left-0">
@@ -1,7 +1,12 @@
import { twMerge } from "tailwind-merge";
import { TabSelector, TabItem } from "~/components/ui/TabSelector";
import { useComputed } from "@preact/signals-react";
import { appMode, dispatchUiEvents } from "~/signals";
import {
appMode,
dispatchUiEvents,
GenerationLoadingState,
generationSignal,
} from "~/signals";
// import { paperWrapperStyles } from "~/components/styles";
// import { faPlus, faQuestion } from "@fortawesome/pro-solid-svg-icons";
@@ -12,21 +17,26 @@ export const ToolbarTopLeft = () => {
const activeTab = useComputed(() => appMode.value);
const tabs: TabItem[] = [
{ id: "image", label: "Image" },
{ id: "realtime", label: "Realtime" },
{ id: "edit", label: "Edit" },
{ id: "generate", label: "Generate" },
// { id: "video", label: "Video" },
];
const handleTabChange = (tabId: string) => {
// Update the global signal
dispatchUiEvents.changeAppMode(tabId as "image" | "edit");
dispatchUiEvents.changeAppMode(tabId as "realtime" | "edit" | "generate");
};
const generationState = generationSignal.value;
console.log(appMode.value);
return (
<div
className={twMerge("z-20 flex h-fit w-fit items-center gap-8 pl-1 pr-4")}
className={twMerge(
"relative z-50 flex h-fit w-fit items-center gap-8 pl-1 pr-4",
)}
>
<img
src="/brand/mira-logo.png"
@@ -38,6 +48,9 @@ export const ToolbarTopLeft = () => {
tabs={tabs}
activeTab={activeTab.value}
onTabChange={handleTabChange}
disabled={
generationState.loadingState === GenerationLoadingState.GENERATING
}
/>
{/* <ToolbarButton icon={faPlus}>
@@ -22,7 +22,9 @@ export const ToolbarTopRight = () => {
const [aspectRatio, setAspectRatio] = useState("1:1");
return (
<div className={twMerge("z-20 flex h-fit w-fit items-center gap-2")}>
<div
className={twMerge("relative z-50 flex h-fit w-fit items-center gap-2")}
>
{import.meta.env.DEV && ( // only show in dev mode for now
<ToggleButton
label="Ratio"
@@ -1,6 +1,7 @@
import React, { useRef, useEffect, useState } from "react";
import { Tab, TabGroup, TabList } from "@headlessui/react";
import { twMerge } from "tailwind-merge";
import { Tooltip } from "./Tooltip";
export interface TabItem {
id: string;
@@ -12,6 +13,8 @@ interface TabSelectorProps {
activeTab: string;
onTabChange: (tabId: string) => void;
className?: string;
disabled?: boolean;
disabledMessage?: string;
}
export const TabSelector: React.FC<TabSelectorProps> = ({
@@ -19,6 +22,8 @@ export const TabSelector: React.FC<TabSelectorProps> = ({
activeTab,
onTabChange,
className,
disabled,
disabledMessage,
}) => {
// Find the index of the active tab
const selectedIndex = tabs.findIndex((tab) => tab.id === activeTab);
@@ -46,13 +51,17 @@ export const TabSelector: React.FC<TabSelectorProps> = ({
}
}, [selectedIndex, tabs]);
return (
const TabGroupElement = (
<TabGroup
selectedIndex={selectedIndex}
onChange={handleTabChange}
className={twMerge("w-full", className)}
className={twMerge(
"w-full",
className,
disabled ? "cursor-not-allowed opacity-60" : "",
)}
>
<TabList className="glass relative inline-flex min-w-fit overflow-x-auto rounded-lg p-1 shadow-lg">
<TabList className="glass relative inline-flex min-w-fit overflow-x-auto rounded-lg p-0.5 py-1 shadow-lg">
{/* Animated indicator */}
<div
className="absolute top-1 z-10 h-[calc(100%-8px)] rounded-md border-2 border-primary bg-primary/30 transition-all duration-200 ease-in-out"
@@ -66,10 +75,12 @@ export const TabSelector: React.FC<TabSelectorProps> = ({
<Tab
key={tab.id}
ref={(el) => (tabsRef.current[index] = el)}
disabled={disabled}
className={({ selected }) =>
twMerge(
"relative z-20 mx-0.5 min-w-max rounded-md border-2 border-transparent px-4 py-0.5 text-center font-semibold transition-all duration-200 ease-in-out",
selected ? "text-white" : "text-gray-300 hover:text-white",
disabled ? "cursor-not-allowed opacity-60" : "",
)
}
>
@@ -79,4 +90,18 @@ export const TabSelector: React.FC<TabSelectorProps> = ({
</TabList>
</TabGroup>
);
return (
<>
{disabled && (
<Tooltip
tip={disabledMessage ?? "Cannot change tab. Generation in progress."}
position="bottom"
>
{TabGroupElement}
</Tooltip>
)}
{!disabled && TabGroupElement}
</>
);
};
@@ -0,0 +1,50 @@
import { effect, signal } from "@preact/signals-react";
export enum GenerationLoadingState {
INIT,
GENERATING,
GENERATED,
}
const initValue = { loadingState: GenerationLoadingState.INIT, prompt: "" };
export const generationSignal = signal<{
loadingState: GenerationLoadingState;
prompt: string;
imageB64?: string;
}>(initValue);
export const startGeneration = (prompt: string) => {
generationSignal.value = {
loadingState: GenerationLoadingState.GENERATING,
prompt,
imageB64: undefined,
};
};
export const onStartGeneration = (callback: (prompt: string) => void) => {
effect(() => {
// If value just changed to GENERATING, notify the listeners
if (
generationSignal.value.loadingState === GenerationLoadingState.GENERATING
) {
callback(generationSignal.value.prompt);
}
});
};
export const finishGeneration = (imageB64: string) => {
generationSignal.value = {
loadingState: GenerationLoadingState.GENERATED,
prompt: generationSignal.value.prompt,
imageB64,
};
};
export const cancelGeneration = () => {
if (
generationSignal.value.loadingState === GenerationLoadingState.GENERATING
) {
generationSignal.value = initValue;
}
};
@@ -0,0 +1 @@
export * from "./generationSignals";
+1
View File
@@ -1,3 +1,4 @@
export * from "./authentication";
export * from "./uiAccess";
export * from "./uiEvents";
export * from "./generation";
@@ -3,10 +3,10 @@ import { signal, effect } from "@preact/signals-core";
// this is to control the app mode - BFlat
// Define the app mode types
export type AppModeType = "image" | "edit";
export type AppModeType = "realtime" | "edit" | "generate";
// Create a signal for the current app mode
export const appMode = signal<AppModeType>("image");
export const appMode = signal<AppModeType>("realtime");
// Create a function to change the app mode
export const changeAppMode = (mode: AppModeType) => {
+25 -17
View File
@@ -200,6 +200,14 @@
"@babel/helper-validator-identifier" "^7.24.7"
to-fast-properties "^2.0.0"
"@devnomic/marquee@^1.0.3":
version "1.0.3"
resolved "https://registry.npmjs.org/@devnomic/marquee/-/marquee-1.0.3.tgz"
integrity sha512-q8QUtNO8ff+UvfrKJsHBtxh4fbOfhGk6SSdK0Uu/XZFdEzrmiGJ/l0FuVHRYsPAmzgnhVVnz1dGNB5CWGLMVvw==
dependencies:
clsx "^2.1.0"
tailwind-merge "^2.2.0"
"@emotion/is-prop-valid@^1.1.0":
version "1.3.0"
resolved "https://registry.npmjs.org/@emotion/is-prop-valid/-/is-prop-valid-1.3.0.tgz"
@@ -222,11 +230,6 @@
resolved "https://registry.npmjs.org/@emotion/unitless/-/unitless-0.7.5.tgz"
integrity sha512-OWORNpfjMsSSUBVrRBVGECkhWcULOAJz9ZW8uK9qgxD+87M7jHRcvh/A96XXNhXTLmKcoYSQtBEX7lHMO7YRwg==
"@esbuild/win32-x64@0.21.5":
version "0.21.5"
resolved "https://registry.npmjs.org/@esbuild/win32-x64/-/win32-x64-0.21.5.tgz"
integrity sha512-tQd/1efJuzPC6rCFwEvLtci/xNFcTZknmXs98FYDfGE4wP9ClFV98nyKrzJKVPMhdDnjzLhdUyMX4PsQAPjwIw==
"@eslint-community/eslint-utils@^4.2.0", "@eslint-community/eslint-utils@^4.4.0":
version "4.4.0"
resolved "https://registry.npmjs.org/@eslint-community/eslint-utils/-/eslint-utils-4.4.0.tgz"
@@ -741,10 +744,10 @@
resolved "https://registry.npmjs.org/@remix-run/router/-/router-1.19.1.tgz"
integrity sha512-S45oynt/WH19bHbIXjtli6QmwNYvaz+vtnubvNpNDvUOoA/OWh6j1OikIP3G+v5GHdxyC6EXoChG3HgYGEUfcg==
"@rollup/rollup-win32-x64-msvc@4.24.0":
"@rollup/rollup-darwin-arm64@4.24.0":
version "4.24.0"
resolved "https://registry.npmjs.org/@rollup/rollup-win32-x64-msvc/-/rollup-win32-x64-msvc-4.24.0.tgz"
integrity sha512-fbMkAF7fufku0N2dE5TBXcNlg0pt0cJue4xBRE2Qc5Vqikxr4VCgKj/ht6SMdFcOacVA9rqF70APJ8RN/4vMJw==
resolved "https://registry.npmjs.org/@rollup/rollup-darwin-arm64/-/rollup-darwin-arm64-4.24.0.tgz"
integrity sha512-bIv+X9xeSs1XCk6DVvkO+S/z8/2AMt/2lMqdQbMrmVpgFvXlmde9mLcbQpztXm1tajC3raFDqegsH18HQPMYtA==
"@sigstore/bundle@^2.3.2":
version "2.3.2"
@@ -816,10 +819,10 @@
resolved "https://registry.npmjs.org/@tauri-apps/api/-/api-2.2.0.tgz"
integrity sha512-R8epOeZl1eJEl603aUMIGb4RXlhPjpgxbGVEaqY+0G5JG9vzV/clNlzTeqc+NLYXVqXcn8mb4c5b9pJIUDEyAg==
"@tauri-apps/cli-win32-x64-msvc@2.2.7":
"@tauri-apps/cli-darwin-arm64@2.2.7":
version "2.2.7"
resolved "https://registry.npmjs.org/@tauri-apps/cli-win32-x64-msvc/-/cli-win32-x64-msvc-2.2.7.tgz"
integrity sha512-ZB8Kw90j8Ld+9tCWyD2fWCYfIrzbQohJ4DJSidNwbnehlZzP7wAz6Z3xjsvUdKtQ3ibtfoeTqVInzCCEpI+pWg==
resolved "https://registry.npmjs.org/@tauri-apps/cli-darwin-arm64/-/cli-darwin-arm64-2.2.7.tgz"
integrity sha512-54kcpxZ3X1Rq+pPTzk3iIcjEVY4yv493uRx/80rLoAA95vAC0c//31Whz75UVddDjJfZvXlXZ3uSZ+bnCOnt0A==
"@tauri-apps/cli@^2.2.7":
version "2.2.7"
@@ -910,14 +913,14 @@
resolved "https://registry.npmjs.org/@types/prop-types/-/prop-types-15.7.12.tgz"
integrity sha512-5zvhXYtRNRluoE/jAp4GVsSduVUzNWKkOZrCDBWYtE7biZywwdC2AcEzg+cSMLFRfVgeAFqpfNabiPjxFddV1Q==
"@types/react-dom@^18.3.0":
"@types/react-dom@^18.2.18", "@types/react-dom@^18.3.0":
version "18.3.0"
resolved "https://registry.npmjs.org/@types/react-dom/-/react-dom-18.3.0.tgz"
integrity sha512-EhwApuTmMBmXuFOikhQLIBUn6uFg81SwLMOAUgodJF14SOBOCMdU04gDoYi0WOJJHD144TL32z4yDqCW3dnkQg==
dependencies:
"@types/react" "*"
"@types/react@*", "@types/react@^18.3.3":
"@types/react@*", "@types/react@^18.2.47", "@types/react@^18.3.3":
version "18.3.4"
resolved "https://registry.npmjs.org/@types/react/-/react-18.3.4.tgz"
integrity sha512-J7W30FTdfCxDDjmfRM+/JqLHBIyl7xUIp9kwK637FGmY7+mkSFSe6L4jpZzhj5QMfLssSDP4/i75AKkrdC7/Jw==
@@ -1416,7 +1419,7 @@ cli-columns@^4.0.0:
string-width "^4.2.3"
strip-ansi "^6.0.1"
clsx@^2.0.0:
clsx@^2.0.0, clsx@^2.1.0:
version "2.1.1"
resolved "https://registry.npmjs.org/clsx/-/clsx-2.1.1.tgz"
integrity sha512-eYm0QWBtUrBWZWG0d386OGAw16Z995PiOVo2B7bjWSbHedGl5e0ZWaq65kOGgUSNesEIDkB9ISbTg/JK9dhCZA==
@@ -2082,6 +2085,11 @@ fs-minipass@^3.0.0, fs-minipass@^3.0.3:
dependencies:
minipass "^7.0.3"
fsevents@~2.3.2, fsevents@~2.3.3:
version "2.3.3"
resolved "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz"
integrity sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==
function-bind@^1.1.2:
version "1.1.2"
resolved "https://registry.npmjs.org/function-bind/-/function-bind-1.1.2.tgz"
@@ -3635,7 +3643,7 @@ react-colorful@^5.6.1:
resolved "https://registry.npmjs.org/react-colorful/-/react-colorful-5.6.1.tgz"
integrity sha512-1exovf0uGTGyq5mXQT0zgQ80uvj2PCwvF8zY1RN9/vbJVSjSo3fsB/4L3ObbF7u70NduSiK4xu4Y6q1MHoUGEw==
"react-dom@^16.8.0 || ^17.0.0 || ^18.0.0", react-dom@^18, react-dom@^18.0.0, "react-dom@^18.0.0 || ^19.0.0 || ^19.0.0-rc", react-dom@^18.3.1, "react-dom@>= 16.8.0", react-dom@>=16.4.0, react-dom@>=16.8, react-dom@>=16.8.0:
"react-dom@^16.8.0 || ^17.0.0 || ^18.0.0", react-dom@^18, react-dom@^18.0.0, "react-dom@^18.0.0 || ^19.0.0 || ^19.0.0-rc", react-dom@^18.2.0, react-dom@^18.3.1, "react-dom@>= 16.8.0", react-dom@>=16.4.0, react-dom@>=16.8, react-dom@>=16.8.0:
version "18.3.1"
resolved "https://registry.npmjs.org/react-dom/-/react-dom-18.3.1.tgz"
integrity sha512-5m4nQKp+rZRb09LNH59GM4BxTh9251/ylbKIbpe7TpGxfJ+9kv6BLkLBXIjjspbgbnIBNqlI23tRnTWT0snUIw==
@@ -3684,7 +3692,7 @@ react-router@6.26.1:
dependencies:
"@remix-run/router" "1.19.1"
"react@^16.14.0 || 17.x || 18.x", "react@^16.8.0 || ^17.0.0 || ^18.0.0", "react@^16.8.0 || ^17.0.0-rc.1 || ^18.0.0 || ^19.0.0", react@^18, react@^18.0.0, "react@^18.0.0 || ^19.0.0 || ^19.0.0-rc", react@^18.3.1, "react@>= 16.8.0", react@>=16.3, react@>=16.4.0, react@>=16.8, react@>=16.8.0:
"react@^16.14.0 || 17.x || 18.x", "react@^16.8.0 || ^17.0.0 || ^18.0.0", "react@^16.8.0 || ^17.0.0-rc.1 || ^18.0.0 || ^19.0.0", react@^18, react@^18.0.0, "react@^18.0.0 || ^19.0.0 || ^19.0.0-rc", react@^18.2.0, react@^18.3.1, "react@>= 16.8.0", react@>=16.3, react@>=16.4.0, react@>=16.8, react@>=16.8.0:
version "18.3.1"
resolved "https://registry.npmjs.org/react/-/react-18.3.1.tgz"
integrity sha512-wS+hAgJShR0KhEvPJArfuPVN1+Hz1t0Y6n5jLrGQbkb4urgPE/0Rve+1kMB1v/oWgHgm4WIcV+i7F2pTVj+2iQ==
@@ -4180,7 +4188,7 @@ tabbable@^6.0.0:
resolved "https://registry.npmjs.org/tabbable/-/tabbable-6.2.0.tgz"
integrity sha512-Cat63mxsVJlzYvN51JmVXIgNoUokrIaT2zLclCXjRd8boZ0004U4KCs/sToJ75C6sdlByWxpYnb5Boif1VSFew==
tailwind-merge@^2.5.2:
tailwind-merge@^2.2.0, tailwind-merge@^2.5.2:
version "2.5.2"
resolved "https://registry.npmjs.org/tailwind-merge/-/tailwind-merge-2.5.2.tgz"
integrity sha512-kjEBm+pvD+6eAwzJL2Bi+02/9LFLal1Gs61+QB7HvTfQQ0aXwC5LGT8PEt1gS0CWKktKe6ysPTAy3cBC5MeiIg==