mirror of
https://github.com/storytold/artcraft.git
synced 2026-11-02 08:28:18 +00:00
Merge branch 'main' of github.com:storytold/storyteller-rust
This commit is contained in:
CrossPr0duct committed
commit
0233105393
33 files changed
+2231
-1599
No files matched your search
Generated
+53
-37
@@ -672,6 +672,7 @@ dependencies = [
|
||||
"memory_caching",
|
||||
"once_cell",
|
||||
"rand 0.8.5",
|
||||
"reqwest 0.12.12",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_yaml",
|
||||
@@ -1079,7 +1080,7 @@ dependencies = [
|
||||
"rustversion",
|
||||
"serde",
|
||||
"sync_wrapper 0.1.2",
|
||||
"tower",
|
||||
"tower 0.4.13",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
]
|
||||
@@ -3316,7 +3317,7 @@ dependencies = [
|
||||
"google_drive_common",
|
||||
"hashing",
|
||||
"hostname 0.4.0",
|
||||
"hyper 1.3.1",
|
||||
"hyper 1.6.0",
|
||||
"jobs_common",
|
||||
"log",
|
||||
"mimetypes",
|
||||
@@ -4762,7 +4763,7 @@ dependencies = [
|
||||
"errors",
|
||||
"jwt-simple",
|
||||
"my-workspace-hack",
|
||||
"reqwest 0.12.7",
|
||||
"reqwest 0.12.12",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
"serde_json",
|
||||
@@ -4988,7 +4989,7 @@ dependencies = [
|
||||
"native-tls",
|
||||
"num_cpus",
|
||||
"rand 0.8.5",
|
||||
"reqwest 0.12.7",
|
||||
"reqwest 0.12.12",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.11",
|
||||
@@ -5192,9 +5193,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "httparse"
|
||||
version = "1.8.0"
|
||||
version = "1.10.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d897f394bad6a705d5f4104762e116a75639e470d80901eed05a860a95cb1904"
|
||||
checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
|
||||
|
||||
[[package]]
|
||||
name = "httpdate"
|
||||
@@ -5234,9 +5235,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "hyper"
|
||||
version = "1.3.1"
|
||||
version = "1.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fe575dd17d0862a9a33781c8c4696a55c320909004a67a00fb286ba8b1bc496d"
|
||||
checksum = "cc2b571658e38e0c01b1fdca3bbbe93c00d3d71693ff2770043f8c29bc7d6f80"
|
||||
dependencies = [
|
||||
"bytes 1.10.0",
|
||||
"futures-channel",
|
||||
@@ -5260,7 +5261,7 @@ checksum = "08afdbb5c31130e3034af566421053ab03787c640246a446327f550d11bcb333"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"http 1.1.0",
|
||||
"hyper 1.3.1",
|
||||
"hyper 1.6.0",
|
||||
"hyper-util",
|
||||
"rustls 0.23.13",
|
||||
"rustls-pki-types",
|
||||
@@ -5303,7 +5304,7 @@ checksum = "70206fc6890eaca9fde8a0bf71caa2ddfc9fe045ac9e5c70df101a7dbde866e0"
|
||||
dependencies = [
|
||||
"bytes 1.10.0",
|
||||
"http-body-util",
|
||||
"hyper 1.3.1",
|
||||
"hyper 1.6.0",
|
||||
"hyper-util",
|
||||
"native-tls",
|
||||
"tokio",
|
||||
@@ -5313,20 +5314,19 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "hyper-util"
|
||||
version = "0.1.5"
|
||||
version = "0.1.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7b875924a60b96e5d7b9ae7b066540b1dd1cbd90d1828f54c92e02a283351c56"
|
||||
checksum = "df2dcfbe0677734ab2f3ffa7fa7bfd4706bfdc1ef393f2ee30184aed67e631b4"
|
||||
dependencies = [
|
||||
"bytes 1.10.0",
|
||||
"futures-channel",
|
||||
"futures-util",
|
||||
"http 1.1.0",
|
||||
"http-body 1.0.0",
|
||||
"hyper 1.3.1",
|
||||
"hyper 1.6.0",
|
||||
"pin-project-lite",
|
||||
"socket2",
|
||||
"tokio",
|
||||
"tower",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
@@ -5563,7 +5563,7 @@ dependencies = [
|
||||
"rand 0.8.5",
|
||||
"redis_common",
|
||||
"redis_schema",
|
||||
"reqwest 0.12.7",
|
||||
"reqwest 0.12.12",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
"serde_json",
|
||||
@@ -6852,7 +6852,7 @@ dependencies = [
|
||||
"heck 0.4.1",
|
||||
"hmac 0.12.1",
|
||||
"hyper 0.14.28",
|
||||
"hyper 1.3.1",
|
||||
"hyper 1.6.0",
|
||||
"hyper-util",
|
||||
"indexmap 2.2.6",
|
||||
"itertools 0.12.1",
|
||||
@@ -6878,7 +6878,7 @@ dependencies = [
|
||||
"regex",
|
||||
"regex-automata",
|
||||
"regex-syntax",
|
||||
"reqwest 0.12.7",
|
||||
"reqwest 0.12.12",
|
||||
"rsa",
|
||||
"rustls-pki-types",
|
||||
"semver 1.0.23",
|
||||
@@ -6904,7 +6904,7 @@ dependencies = [
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tokio-util",
|
||||
"tower",
|
||||
"tower 0.4.13",
|
||||
"tracing",
|
||||
"unicode-bidi",
|
||||
"unicode-normalization",
|
||||
@@ -9269,9 +9269,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "reqwest"
|
||||
version = "0.12.7"
|
||||
version = "0.12.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8f4955649ef5c38cc7f9e8aa41761d48fb9677197daea9984dc54f56aad5e63"
|
||||
checksum = "43e734407157c3c2034e0258f5e4473ddb361b1e85f95a66690d67264d7cd1da"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bytes 1.10.0",
|
||||
@@ -9283,7 +9283,7 @@ dependencies = [
|
||||
"http 1.1.0",
|
||||
"http-body 1.0.0",
|
||||
"http-body-util",
|
||||
"hyper 1.3.1",
|
||||
"hyper 1.6.0",
|
||||
"hyper-rustls",
|
||||
"hyper-tls 0.6.0",
|
||||
"hyper-util",
|
||||
@@ -9308,6 +9308,7 @@ dependencies = [
|
||||
"tokio-native-tls",
|
||||
"tokio-rustls",
|
||||
"tokio-util",
|
||||
"tower 0.5.2",
|
||||
"tower-service",
|
||||
"url",
|
||||
"wasm-bindgen",
|
||||
@@ -9502,7 +9503,7 @@ dependencies = [
|
||||
"minidom",
|
||||
"percent-encoding",
|
||||
"quick-xml 0.36.2",
|
||||
"reqwest 0.12.7",
|
||||
"reqwest 0.12.12",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
"serde_json",
|
||||
@@ -10833,7 +10834,7 @@ dependencies = [
|
||||
"redis_common",
|
||||
"redis_schema",
|
||||
"regex",
|
||||
"reqwest 0.12.7",
|
||||
"reqwest 0.12.12",
|
||||
"reusable_types",
|
||||
"rust-s3",
|
||||
"serde",
|
||||
@@ -11597,7 +11598,7 @@ dependencies = [
|
||||
"percent-encoding",
|
||||
"plist",
|
||||
"raw-window-handle",
|
||||
"reqwest 0.12.7",
|
||||
"reqwest 0.12.12",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_repr",
|
||||
@@ -11945,7 +11946,7 @@ dependencies = [
|
||||
"easyenv",
|
||||
"log",
|
||||
"my-workspace-hack",
|
||||
"reqwest 0.12.7",
|
||||
"reqwest 0.12.12",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
@@ -12399,7 +12400,7 @@ dependencies = [
|
||||
"prost",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tower",
|
||||
"tower 0.4.13",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
@@ -12435,16 +12436,31 @@ dependencies = [
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tower-layer"
|
||||
version = "0.3.2"
|
||||
name = "tower"
|
||||
version = "0.5.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c20c8dbed6283a09604c3e69b4b7eeb54e298b8a600d4d5ecb5ad39de609f1d0"
|
||||
checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"pin-project-lite",
|
||||
"sync_wrapper 1.0.1",
|
||||
"tokio",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tower-layer"
|
||||
version = "0.3.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e"
|
||||
|
||||
[[package]]
|
||||
name = "tower-service"
|
||||
version = "0.3.2"
|
||||
version = "0.3.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b6bc1c9ce2b5135ac7f93c72918fc37feb872bdc6a5533a8b85eb4b86bfdae52"
|
||||
checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3"
|
||||
|
||||
[[package]]
|
||||
name = "tracing"
|
||||
@@ -12581,7 +12597,7 @@ dependencies = [
|
||||
"google_drive_common",
|
||||
"hashing",
|
||||
"hostname 0.4.0",
|
||||
"hyper 1.3.1",
|
||||
"hyper 1.6.0",
|
||||
"jobs_common",
|
||||
"log",
|
||||
"my-workspace-hack",
|
||||
@@ -12619,7 +12635,7 @@ dependencies = [
|
||||
"my-workspace-hack",
|
||||
"mysql_queries",
|
||||
"r2d2_redis",
|
||||
"reqwest 0.12.7",
|
||||
"reqwest 0.12.12",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
"serde_json",
|
||||
@@ -12979,7 +12995,7 @@ dependencies = [
|
||||
"actix-web",
|
||||
"mime_guess",
|
||||
"regex",
|
||||
"reqwest 0.12.7",
|
||||
"reqwest 0.12.12",
|
||||
"rust-embed",
|
||||
"serde",
|
||||
"serde_json",
|
||||
@@ -13108,7 +13124,7 @@ dependencies = [
|
||||
"google_drive_common",
|
||||
"hashing",
|
||||
"hostname 0.4.0",
|
||||
"hyper 1.3.1",
|
||||
"hyper 1.6.0",
|
||||
"jobs_common",
|
||||
"log",
|
||||
"my-workspace-hack",
|
||||
@@ -13138,7 +13154,7 @@ dependencies = [
|
||||
"filesys",
|
||||
"google_drive_common",
|
||||
"hostname 0.4.0",
|
||||
"hyper 1.3.1",
|
||||
"hyper 1.6.0",
|
||||
"jobs_common",
|
||||
"log",
|
||||
"my-workspace-hack",
|
||||
@@ -13315,7 +13331,7 @@ dependencies = [
|
||||
"log",
|
||||
"my-workspace-hack",
|
||||
"once_cell",
|
||||
"reqwest 0.12.7",
|
||||
"reqwest 0.12.12",
|
||||
"rss",
|
||||
"scraper",
|
||||
"serde",
|
||||
|
||||
+1
-1
@@ -2,6 +2,6 @@ image_width: 512
|
||||
image_height: 512
|
||||
#scheduler_steps: 15
|
||||
# NB(bt): It looks like we can get down to 2 sampling steps, perhaps ?
|
||||
scheduler_steps: 2
|
||||
scheduler_steps: 4
|
||||
scheduler_samples: 1
|
||||
cfg_scale: 1.1
|
||||
@@ -39,6 +39,7 @@ serde_yaml = "0.9"
|
||||
log = "0.4"
|
||||
once_cell = "1.20.3"
|
||||
rand = "0.8"
|
||||
reqwest = "0.12.12"
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
tauri = { version = "2.2.4", features = [] }
|
||||
@@ -47,7 +48,7 @@ tokenizers = "0.15.2"
|
||||
|
||||
[features]
|
||||
default = ["cuda"]
|
||||
cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
|
||||
cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda", "candle-transformers/flash-attn", "candle-core/cudnn"]
|
||||
|
||||
[dev-dependencies]
|
||||
libloading = "0.8.6"
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
use crate::ml::downloads::download_all_models::download_all_models;
|
||||
use crate::state::app_config::AppConfig;
|
||||
use tauri::{AppHandle, State};
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn download_models(
|
||||
model_config: State<'_, AppConfig>,
|
||||
app: AppHandle,
|
||||
) -> Result<String, String> {
|
||||
|
||||
download_all_models().await
|
||||
.map_err(|err| {
|
||||
"error downloading models".to_string()
|
||||
})?;
|
||||
|
||||
Ok(("downloaded".to_string()))
|
||||
}
|
||||
@@ -1,6 +1,6 @@
|
||||
use crate::ml::model_cache::ModelCache;
|
||||
use crate::ml::prompt_cache::PromptCache;
|
||||
use crate::ml::stable_diffusion::stable_diffusion_pipeline::{stable_diffusion_pipeline, Args};
|
||||
use crate::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
|
||||
use crate::state::app_config::AppConfig;
|
||||
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
|
||||
use bytes::BytesMut;
|
||||
@@ -10,6 +10,7 @@ use log::{error, info};
|
||||
use std::io::Cursor;
|
||||
use std::path::PathBuf;
|
||||
use tauri::{AppHandle, State};
|
||||
use crate::ml::downloads::download_all_models::download_all_models;
|
||||
|
||||
const PROMPT_FILENAME : &str = "prompt.txt";
|
||||
const PROMPT: &str = "A beautiful landscape with mountains and a lake";
|
||||
@@ -40,7 +41,14 @@ pub async fn infer_image(
|
||||
let image = hydrate_base64_image(image)
|
||||
.map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
|
||||
|
||||
let result = do_infer_image(&prompt, image, strength, &model_config, &model_cache, prompt_cache, app);
|
||||
// TODO(bt): Move this to another endpoint
|
||||
download_all_models().await
|
||||
.map_err(|err| {
|
||||
error!("couldn't download models: {:?}", err);
|
||||
"Couldn't download models".to_string()
|
||||
})?;
|
||||
|
||||
let result = do_infer_image(&prompt, image, strength, &model_config, &model_cache, prompt_cache, app).await;
|
||||
|
||||
if let Err(err) = result.as_deref() {
|
||||
error!("There was an error: {:?}", err);
|
||||
@@ -49,22 +57,20 @@ pub async fn infer_image(
|
||||
result
|
||||
}
|
||||
|
||||
fn do_infer_image(
|
||||
async fn do_infer_image(
|
||||
prompt: &str,
|
||||
image: DynamicImage,
|
||||
strength: Option<u8>,
|
||||
config: &AppConfig,
|
||||
model_cache: &ModelCache,
|
||||
prompt_cache: State<PromptCache>,
|
||||
prompt_cache: State<'_, PromptCache>,
|
||||
app: AppHandle,
|
||||
) -> Result<String, String> {
|
||||
println!("infer_image called; generating image with SDXL Turbo...");
|
||||
|
||||
|
||||
let args = Args {
|
||||
image: &image,
|
||||
prompt: prompt.to_string(),
|
||||
uncond_prompt: NEGATIVE_PROMPT.to_string(),
|
||||
//guidance_scale: Some(0.0),
|
||||
model_cache,
|
||||
configs: config,
|
||||
prompt_cache: &prompt_cache,
|
||||
@@ -73,7 +79,7 @@ fn do_infer_image(
|
||||
app: &app,
|
||||
};
|
||||
|
||||
match stable_diffusion_pipeline(args) {
|
||||
match lcm_pipeline(args) {
|
||||
Ok(image) => {
|
||||
let mut bytes = Vec::with_capacity(1024*1024);
|
||||
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod download_models;
|
||||
pub mod flip_image;
|
||||
pub mod image_endpoint;
|
||||
pub mod test_counter;
|
||||
@@ -6,6 +6,7 @@ pub mod state;
|
||||
use crate::ml::model_cache::ModelCache;
|
||||
use crate::ml::prompt_cache::PromptCache;
|
||||
use crate::state::app_config::AppConfig;
|
||||
use endpoints::download_models::download_models;
|
||||
use endpoints::image_endpoint::infer_image;
|
||||
use endpoints::test_counter::test_counter;
|
||||
use tauri_plugin_log::Target;
|
||||
@@ -50,6 +51,7 @@ pub fn run() {
|
||||
.invoke_handler(tauri::generate_handler![
|
||||
test_counter,
|
||||
infer_image,
|
||||
download_models,
|
||||
])
|
||||
.run(tauri::generate_context!())
|
||||
.expect("error while running tauri application");
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
use crate::ml::downloads::download_model_registry_file::download_model_registry_file;
|
||||
use crate::ml::model_registry::ModelRegistry;
|
||||
|
||||
pub async fn download_all_models() -> anyhow::Result<()> {
|
||||
download_model_registry_file(ModelRegistry::ClipJson).await?;
|
||||
download_model_registry_file(ModelRegistry::SdxlTurboUnet).await?;
|
||||
download_model_registry_file(ModelRegistry::SdxlTurboVae).await?;
|
||||
download_model_registry_file(ModelRegistry::SdxlTurboClipEncoder).await?;
|
||||
download_model_registry_file(ModelRegistry::SdxlTurboClipEncoder2).await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
use std::fs::File;
|
||||
use std::path::{Path, PathBuf};
|
||||
use log::info;
|
||||
use crate::ml::model_registry::ModelRegistry;
|
||||
|
||||
pub async fn download_model_registry_file(model_type: ModelRegistry) -> anyhow::Result<PathBuf> {
|
||||
info!("downloading model: {:?}", model_type);
|
||||
let path = PathBuf::from(model_type.get_filename());
|
||||
if path.exists() {
|
||||
info!("model already exists: {:?}", path);
|
||||
return Ok(path);
|
||||
}
|
||||
|
||||
let body = reqwest::get(model_type.get_download_url())
|
||||
.await?
|
||||
.bytes()
|
||||
.await?;
|
||||
|
||||
let _file = File::create(&path)?;
|
||||
|
||||
std::fs::write(&path, body.as_ref())?;
|
||||
|
||||
Ok(path)
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod download_model_registry_file;
|
||||
pub mod download_all_models;
|
||||
@@ -3,4 +3,6 @@ pub mod model_cache;
|
||||
pub mod model_file;
|
||||
pub mod models;
|
||||
pub mod prompt_cache;
|
||||
pub mod stable_diffusion;
|
||||
pub mod stable_diffusion;
|
||||
pub mod model_registry;
|
||||
pub mod downloads;
|
||||
@@ -0,0 +1,41 @@
|
||||
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub enum ModelRegistry {
|
||||
ClipJson,
|
||||
SdxlTurboUnet,
|
||||
SdxlTurboVae,
|
||||
SdxlTurboClipEncoder,
|
||||
SdxlTurboClipEncoder2,
|
||||
}
|
||||
|
||||
impl ModelRegistry {
|
||||
pub fn get_hf_id(&self) -> Option<&'static str> {
|
||||
match self {
|
||||
Self::ClipJson => Some("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"),
|
||||
Self::SdxlTurboUnet
|
||||
| Self::SdxlTurboVae
|
||||
| Self::SdxlTurboClipEncoder
|
||||
| Self::SdxlTurboClipEncoder2 => Some("stabilityai/sdxl-turbo"),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_download_url(&self) -> &'static str {
|
||||
match self {
|
||||
Self::ClipJson => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/tokenizer.json",
|
||||
Self::SdxlTurboUnet => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/diffusion_pytorch_model.unet.safetensors",
|
||||
Self::SdxlTurboVae => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/diffusion_pytorch_model.vae.safetensors",
|
||||
Self::SdxlTurboClipEncoder => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/clip_text_encoder.safetensors",
|
||||
Self::SdxlTurboClipEncoder2 => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/clip_text_encoder_2.safetensors",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_filename(&self) -> &'static str {
|
||||
match self {
|
||||
Self::ClipJson => "clip_vit_bigg_14_lion2b_39b_b160k.tokenizer.json",
|
||||
Self::SdxlTurboUnet => "diffusion_pytorch_model.unet.safetensors",
|
||||
Self::SdxlTurboVae => "diffusion_pytorch_model.vae.safetensors",
|
||||
Self::SdxlTurboClipEncoder => "clip_text_encoder.safetensors",
|
||||
Self::SdxlTurboClipEncoder2 => "clip_text_encoder_2.safetensors",
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -8,28 +8,19 @@ pub struct UNetModel {
|
||||
}
|
||||
|
||||
impl UNetModel {
|
||||
pub fn new<P: AsRef<Path>>(
|
||||
sd_config: &StableDiffusionConfig,
|
||||
unet_weights_file: P,
|
||||
device: &Device,
|
||||
dtype: DType,
|
||||
) -> anyhow::Result<Self> {
|
||||
pub fn new<P: AsRef<Path>>(sd_config: &StableDiffusionConfig, unet_weights_file: P, device: &Device, dtype: DType) -> anyhow::Result<Self> {
|
||||
println!("building unet model... (1)");
|
||||
|
||||
let model = sd_config.build_unet(
|
||||
unet_weights_file,
|
||||
&device,
|
||||
4,
|
||||
false,
|
||||
dtype
|
||||
)?;
|
||||
|
||||
Ok(Self {
|
||||
model,
|
||||
})
|
||||
|
||||
let model = sd_config.build_unet(unet_weights_file, &device, 4, true, dtype)?;
|
||||
|
||||
Ok(Self { model })
|
||||
}
|
||||
|
||||
|
||||
pub fn inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result<Tensor> {
|
||||
Ok(self.model.forward(&latent_model_input, timestep, &text_embeddings)?)
|
||||
}
|
||||
|
||||
pub fn forward_with_guidance(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor, guidance_scale_embedding: &Tensor) -> anyhow::Result<Tensor> {
|
||||
Ok(self.model.forward_with_guidance(&latent_model_input, timestep, &text_embeddings, Some(guidance_scale_embedding))?)
|
||||
}
|
||||
}
|
||||
@@ -3,6 +3,7 @@ extern crate accelerate_src;
|
||||
#[cfg(feature = "mkl")]
|
||||
extern crate intel_mkl_src;
|
||||
|
||||
use std::path::PathBuf;
|
||||
use candle_transformers::models::stable_diffusion;
|
||||
|
||||
use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
|
||||
@@ -10,6 +11,7 @@ use anyhow::Error as E;
|
||||
use candle_core::{DType, Device, Module, Tensor, D};
|
||||
use log::info;
|
||||
use tokenizers::Tokenizer;
|
||||
use crate::ml::model_registry::ModelRegistry;
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn infer_clip_text_embeddings(
|
||||
@@ -76,12 +78,21 @@ fn do_infer_clip_text_embeddings(
|
||||
) -> anyhow::Result<Tensor> {
|
||||
|
||||
let tokenizer_file = if first {
|
||||
info!("ModelFile::Tokenizer");
|
||||
ModelFile::Tokenizer
|
||||
} else {
|
||||
info!("ModelFile::Tokenizer2");
|
||||
ModelFile::Tokenizer2
|
||||
};
|
||||
let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
|
||||
|
||||
info!("Loading clip from download");
|
||||
//let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
|
||||
let tokenizer = PathBuf::from(ModelRegistry::ClipJson.get_filename());
|
||||
|
||||
info!("Tokenizer path: {:?}", tokenizer);
|
||||
|
||||
let tokenizer = Tokenizer::from_file(tokenizer).map_err(E::msg)?;
|
||||
|
||||
let pad_id = match &sd_config.clip.pad_with {
|
||||
Some(padding) => *tokenizer.get_vocab(true).get(padding.as_str()).unwrap(),
|
||||
None => *tokenizer.get_vocab(true).get("<|endoftext|>").unwrap(),
|
||||
|
||||
@@ -0,0 +1,294 @@
|
||||
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
|
||||
|
||||
use crate::events::notification_event::{ModelType, NotificationEvent};
|
||||
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::model_registry::ModelRegistry;
|
||||
use crate::ml::models::unet_model::UNetModel;
|
||||
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::state::app_config::AppConfig;
|
||||
use anyhow::{anyhow, Error as E, Result};
|
||||
use candle_core::{DType, IndexOp, Tensor, D};
|
||||
use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
|
||||
use candle_transformers::models::stable_diffusion::lcm::LCMScheduler;
|
||||
use image::DynamicImage;
|
||||
use log::info;
|
||||
use rand::Rng;
|
||||
use tauri::{AppHandle, Emitter};
|
||||
|
||||
pub struct Args<'a> {
|
||||
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: &'a AppHandle,
|
||||
}
|
||||
|
||||
pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
|
||||
let Args {
|
||||
prompt,
|
||||
uncond_prompt,
|
||||
cfg_scale,
|
||||
i2i_strength,
|
||||
configs,
|
||||
model_cache,
|
||||
prompt_cache,
|
||||
app,
|
||||
image,
|
||||
} = args;
|
||||
|
||||
println!("Starting image generation with the following configuration:");
|
||||
println!(" Model: {:?}", configs.sd_version);
|
||||
println!(" Prompt: {}", prompt);
|
||||
println!(" Steps: {}", configs.scheduler_steps);
|
||||
println!(" Device: {:?}", configs.device);
|
||||
|
||||
println!("Model dimensions: {}x{}", configs.sd_config.width, configs.sd_config.height);
|
||||
|
||||
// Use LCM Scheduler instead of Euler Ancestral for better speed and quality
|
||||
let mut scheduler = LCMScheduler::new(
|
||||
configs.scheduler_steps,
|
||||
candle_transformers::models::stable_diffusion::lcm::LCMSchedulerConfig::default(),
|
||||
)?;
|
||||
|
||||
let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
|
||||
configs.device.set_seed(seed)?;
|
||||
|
||||
info!("Using seed: {}", seed);
|
||||
|
||||
let guidance_scale = match cfg_scale {
|
||||
Some(guidance_scale) => guidance_scale,
|
||||
None => match configs.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"
|
||||
},
|
||||
};
|
||||
|
||||
let use_guide_scale = guidance_scale > 1.0;
|
||||
|
||||
info!("Using guide scale: {}", use_guide_scale);
|
||||
|
||||
info!("Checking if prompt is cached");
|
||||
let maybe_cached = prompt_cache.get_copy(&prompt)?;
|
||||
|
||||
let mut text_embeddings = if let Some(tensor) = maybe_cached {
|
||||
tensor
|
||||
} 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,
|
||||
use_guide_scale,
|
||||
)?;
|
||||
prompt_cache.store_copy(&prompt, &tensor)?;
|
||||
tensor
|
||||
};
|
||||
|
||||
println!("Text embeddings shape: {:?}", text_embeddings.shape());
|
||||
|
||||
println!("Loading input image into tensor...");
|
||||
|
||||
let input_image = dynamic_image_to_tensor(
|
||||
image,
|
||||
&configs.device,
|
||||
configs.dtype)?;
|
||||
|
||||
println!("Reference image shape: {:?}", input_image.shape());
|
||||
|
||||
let vae_scale = get_vae_scale(configs.sd_version);
|
||||
|
||||
let maybe_vae = model_cache.get_vae()?;
|
||||
|
||||
let vae = match maybe_vae {
|
||||
Some(vae) => vae,
|
||||
None => {
|
||||
info!("No vae found in cache; loading...");
|
||||
|
||||
let repo = configs.sd_version.repo();
|
||||
|
||||
println!("Building VAE model from : {:?} ... (3)", repo);
|
||||
|
||||
let vae_file = configs.hf_api.model(repo.to_string())
|
||||
.get("vae/diffusion_pytorch_model.safetensors")?;
|
||||
|
||||
println!("Building VAE model from file {:?}...", &vae_file);
|
||||
|
||||
let mut notify_download_complete = false;
|
||||
//if !vae_file.exists() {
|
||||
if true {
|
||||
notify_download_complete = true;
|
||||
app.emit("notification", NotificationEvent::ModelDownloadStarted {
|
||||
model_name: repo,
|
||||
model_type: ModelType::Vae,
|
||||
})?;
|
||||
}
|
||||
|
||||
let vae = configs
|
||||
.sd_config
|
||||
.build_vae(vae_file, &configs.device, configs.dtype)?;
|
||||
|
||||
let vae = Arc::new(vae);
|
||||
|
||||
model_cache.set_vae(vae.clone())?;
|
||||
|
||||
if notify_download_complete {
|
||||
app.emit("notification", NotificationEvent::ModelDownloadComplete {
|
||||
model_name: repo,
|
||||
model_type: ModelType::Vae,
|
||||
})?;
|
||||
}
|
||||
|
||||
vae
|
||||
}
|
||||
};
|
||||
|
||||
let maybe_unet = model_cache.get_unet()?;
|
||||
|
||||
let unet = match maybe_unet {
|
||||
Some(unet) => unet,
|
||||
None => {
|
||||
info!("No unet found in cache; loading...");
|
||||
|
||||
let mut notify_download_complete = false;
|
||||
//if !unet_file.exists() {
|
||||
if true {
|
||||
notify_download_complete = true;
|
||||
app.emit("notification", NotificationEvent::ModelDownloadStarted {
|
||||
model_name: "sdxl-turbo-unet",
|
||||
model_type: ModelType::Unet,
|
||||
})?;
|
||||
}
|
||||
|
||||
let unet_file = ModelRegistry::SdxlTurboUnet.get_filename();
|
||||
|
||||
let unet = UNetModel::new(&configs.sd_config, unet_file, &configs.device, configs.dtype)
|
||||
.map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
|
||||
|
||||
let unet = Arc::new(unet);
|
||||
|
||||
model_cache.set_unet(unet.clone())?;
|
||||
|
||||
if notify_download_complete {
|
||||
app.emit("notification", NotificationEvent::ModelDownloadComplete {
|
||||
model_name: "sdxl-turbo-unet",
|
||||
model_type: ModelType::Unet,
|
||||
})?;
|
||||
}
|
||||
|
||||
unet
|
||||
}
|
||||
};
|
||||
|
||||
let init_latent_dist : DiagonalGaussianDistribution = vae.encode(&input_image)?;
|
||||
|
||||
// Parse and calculate the img2img strength
|
||||
let img2img_strength = match i2i_strength {
|
||||
None => 0.75f64,
|
||||
Some(strength) => (strength as f64) / 100.0f64,
|
||||
};
|
||||
|
||||
// Create LCM guidance scale embeddings
|
||||
let embedding_dim = 256;
|
||||
let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(
|
||||
guidance_scale,
|
||||
embedding_dim,
|
||||
&configs.device,
|
||||
configs.dtype
|
||||
)?;
|
||||
|
||||
// Generate latents from input image using the LCM approach
|
||||
println!("Generating latents from input image...");
|
||||
let init_latents = (init_latent_dist.sample()? * vae_scale)?;
|
||||
println!("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 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);
|
||||
let noise = init_latents.randn_like(0f64, 1f64)?;
|
||||
let latents = scheduler.add_noise(&init_latents, noise, t_start)?;
|
||||
|
||||
println!("Latents initialized successfully");
|
||||
println!("Initial noise shape: {:?}", latents.shape());
|
||||
|
||||
println!("Starting diffusion process...");
|
||||
|
||||
let mut latents = latents;
|
||||
let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect();
|
||||
|
||||
for (timestep_index, ×tep) in timesteps.iter().enumerate() {
|
||||
println!("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()
|
||||
};
|
||||
|
||||
let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep);
|
||||
|
||||
// Use the guidance scale embedding in the UNet inference with the proper method name
|
||||
let mut noise_pred = match unet.forward_with_guidance(
|
||||
&latent_model_input,
|
||||
timestep as f64,
|
||||
&text_embeddings,
|
||||
&guidance_scale_embedding) {
|
||||
Ok(pred) => pred,
|
||||
Err(e) => {
|
||||
println!("UNet inference failed with error: {}", e);
|
||||
return Err(anyhow::anyhow!("UNet inference failed: {}", e));
|
||||
},
|
||||
};
|
||||
|
||||
if use_guide_scale {
|
||||
let chunks = noise_pred.chunk(2, 0)?;
|
||||
let (noise_pred_uncond, noise_pred_text) = (&chunks[0], &chunks[1]);
|
||||
noise_pred = (noise_pred_uncond + (guidance_scale * (noise_pred_text - noise_pred_uncond)?)?)?;
|
||||
}
|
||||
|
||||
// Apply the LCM scheduler step
|
||||
latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, configs.scheduler_steps)?;
|
||||
}
|
||||
|
||||
println!("Diffusion process completed, decoding image...");
|
||||
let image = vae.decode(&(latents / vae_scale)?)?;
|
||||
|
||||
println!("VAE decode completed");
|
||||
|
||||
println!("Scaling image back");
|
||||
let image = ((image / 2.)? + 0.5)?;
|
||||
|
||||
println!("Normalized image values");
|
||||
let image = (image.clamp(0f32, 1.)? * 255.)?;
|
||||
|
||||
println!("Convert to int8");
|
||||
let image = image.to_dtype(DType::U8)?;
|
||||
|
||||
println!("Converted to 8-bit format");
|
||||
|
||||
let image = tensor_to_image_buffer(&image.i(0)?)?;
|
||||
|
||||
Ok(image)
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod get_vae_scale;
|
||||
pub mod infer_clip_text_embeddings;
|
||||
pub mod stable_diffusion_pipeline;
|
||||
pub mod stable_diffusion_pipeline;
|
||||
pub mod lcm_pipeline;
|
||||
@@ -1,9 +1,12 @@
|
||||
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
|
||||
|
||||
use crate::events::notification_event::{ModelType, NotificationEvent};
|
||||
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::model_registry::ModelRegistry;
|
||||
use crate::ml::models::unet_model::UNetModel;
|
||||
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;
|
||||
@@ -15,8 +18,6 @@ use image::DynamicImage;
|
||||
use log::info;
|
||||
use rand::Rng;
|
||||
use tauri::{AppHandle, Emitter};
|
||||
use crate::events::notification_event::{ModelType, NotificationEvent};
|
||||
use crate::ml::models::unet_model::UNetModel;
|
||||
|
||||
pub struct Args<'a> {
|
||||
pub image: &'a DynamicImage,
|
||||
@@ -122,25 +123,18 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
|
||||
None => {
|
||||
info!("No vae found in cache; loading...");
|
||||
|
||||
let repo = configs.sd_version.repo();
|
||||
|
||||
println!("Building VAE model from : {:?} ... (3)", repo);
|
||||
|
||||
let vae_file = configs.hf_api.model(repo.to_string())
|
||||
.get("vae/diffusion_pytorch_model.safetensors")?;
|
||||
|
||||
println!("Building VAE model from file {:?}...", &vae_file);
|
||||
|
||||
let mut notify_download_complete = false;
|
||||
//if !vae_file.exists() {
|
||||
if true {
|
||||
notify_download_complete = true;
|
||||
app.emit("notification", NotificationEvent::ModelDownloadStarted {
|
||||
model_name: repo,
|
||||
model_name: "sdxl-turbo-vae",
|
||||
model_type: ModelType::Vae,
|
||||
})?;
|
||||
}
|
||||
|
||||
let vae_file = ModelRegistry::SdxlTurboVae.get_filename();
|
||||
|
||||
let vae = configs
|
||||
.sd_config
|
||||
.build_vae(vae_file, &configs.device, configs.dtype)?;
|
||||
@@ -151,7 +145,7 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
|
||||
|
||||
if notify_download_complete {
|
||||
app.emit("notification", NotificationEvent::ModelDownloadComplete {
|
||||
model_name: repo,
|
||||
model_name: "sdxl-turbo-vae",
|
||||
model_type: ModelType::Vae,
|
||||
})?;
|
||||
}
|
||||
@@ -167,23 +161,17 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
|
||||
None => {
|
||||
info!("No unet found in cache; loading...");
|
||||
|
||||
let repo = configs.sd_version.repo();
|
||||
|
||||
info!("Downloading UNET model files from: {} ...", repo);
|
||||
|
||||
let unet_file = configs.hf_api.model(repo.to_string())
|
||||
.get("unet/diffusion_pytorch_model.safetensors")
|
||||
.map_err(|err| anyhow!("error fetching model: {:?}", err))?;
|
||||
|
||||
let mut notify_download_complete = false;
|
||||
//if !unet_file.exists() {
|
||||
if true {
|
||||
notify_download_complete = true;
|
||||
app.emit("notification", NotificationEvent::ModelDownloadStarted {
|
||||
model_name: repo,
|
||||
model_name: "sdxl-turbo-unet",
|
||||
model_type: ModelType::Unet,
|
||||
})?;
|
||||
}
|
||||
|
||||
let unet_file = ModelRegistry::SdxlTurboUnet.get_filename();
|
||||
|
||||
let unet = UNetModel::new(&configs.sd_config, unet_file, &configs.device, configs.dtype)
|
||||
.map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
|
||||
@@ -194,7 +182,7 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
|
||||
|
||||
if notify_download_complete {
|
||||
app.emit("notification", NotificationEvent::ModelDownloadComplete {
|
||||
model_name: repo,
|
||||
model_name: "sdxl-turbo-unet",
|
||||
model_type: ModelType::Unet,
|
||||
})?;
|
||||
}
|
||||
|
||||
@@ -71,7 +71,7 @@ impl AppConfig {
|
||||
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(1),
|
||||
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,
|
||||
|
||||
+932
-738
File diff suppressed because it is too large.
Load diff
@@ -86,6 +86,8 @@ fn main() -> Result<()> {
|
||||
.arg("--expt-relaxed-constexpr")
|
||||
.arg("--expt-extended-lambda")
|
||||
.arg("--use_fast_math")
|
||||
.arg("--compiler-options")
|
||||
.arg("-fPIC")
|
||||
.arg("--verbose");
|
||||
|
||||
if let Ok(target) = std::env::var("TARGET") {
|
||||
|
||||
+1
@@ -32,3 +32,4 @@ metal = ["candle/metal", "dep:candle-metal-kernels", "dep:metal"]
|
||||
[[bench]]
|
||||
name = "bench_main"
|
||||
harness = false
|
||||
|
||||
+453
-695
File diff suppressed because it is too large.
Load diff
@@ -84,6 +84,8 @@ impl StableDiffusionConfig {
|
||||
};
|
||||
// https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/unet/config.json
|
||||
let unet = unet_2d::UNet2DConditionModelConfig {
|
||||
time_embed_dim: None,
|
||||
time_cond_proj_dim: None,
|
||||
blocks: vec![
|
||||
bc(320, Some(1), 8),
|
||||
bc(640, Some(1), 8),
|
||||
@@ -153,6 +155,8 @@ impl StableDiffusionConfig {
|
||||
};
|
||||
// https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/unet/config.json
|
||||
let unet = unet_2d::UNet2DConditionModelConfig {
|
||||
time_cond_proj_dim: None,
|
||||
time_embed_dim: None,
|
||||
blocks: vec![
|
||||
bc(320, Some(1), 5),
|
||||
bc(640, Some(1), 10),
|
||||
@@ -237,6 +241,8 @@ impl StableDiffusionConfig {
|
||||
};
|
||||
// https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/unet/config.json
|
||||
let unet = unet_2d::UNet2DConditionModelConfig {
|
||||
time_cond_proj_dim: None,
|
||||
time_embed_dim: None,
|
||||
blocks: vec![
|
||||
bc(320, None, 5),
|
||||
bc(640, Some(2), 10),
|
||||
@@ -306,6 +312,8 @@ impl StableDiffusionConfig {
|
||||
};
|
||||
// https://huggingface.co/stabilityai/sdxl-turbo/blob/main/unet/config.json
|
||||
let unet = unet_2d::UNet2DConditionModelConfig {
|
||||
time_cond_proj_dim: None,
|
||||
time_embed_dim: None,
|
||||
blocks: vec![
|
||||
bc(320, None, 5),
|
||||
bc(640, Some(2), 10),
|
||||
@@ -405,6 +413,8 @@ impl StableDiffusionConfig {
|
||||
};
|
||||
// https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/unet/config.json
|
||||
let unet = unet_2d::UNet2DConditionModelConfig {
|
||||
time_cond_proj_dim: None,
|
||||
time_embed_dim: None,
|
||||
blocks: vec![
|
||||
bc(320, None, 5),
|
||||
bc(640, Some(2), 10),
|
||||
|
||||
+144
-10
@@ -10,6 +10,10 @@ use candle_nn as nn;
|
||||
use candle_nn::Module;
|
||||
use std::collections::HashMap;
|
||||
|
||||
// fn save_tensor_as_numpy(tensor: &Tensor, filename: &str) -> Result<()> {
|
||||
// // Use Candle's built-in NPY serialization
|
||||
// tensor.write_npy(filename).map_err(|e| anyhow::anyhow!("{}", e))
|
||||
// }
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct BlockConfig {
|
||||
pub out_channels: usize,
|
||||
@@ -33,11 +37,13 @@ pub struct UNet2DConditionModelConfig {
|
||||
pub cross_attention_dim: usize,
|
||||
pub sliced_attention_size: Option<usize>,
|
||||
pub use_linear_projection: bool,
|
||||
pub time_cond_proj_dim: Option<usize>,
|
||||
pub time_embed_dim: Option<usize>,
|
||||
}
|
||||
|
||||
impl Default for UNet2DConditionModelConfig {
|
||||
fn default() -> Self {
|
||||
Self { center_input_sample: false, flip_sin_to_cos: true, freq_shift: 0., blocks: vec![BlockConfig { out_channels: 320, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 640, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: None, attention_head_dim: 8 }], layers_per_block: 2, downsample_padding: 1, mid_block_scale_factor: 1., norm_num_groups: 32, norm_eps: 1e-5, cross_attention_dim: 1280, sliced_attention_size: None, use_linear_projection: false }
|
||||
Self { center_input_sample: false, flip_sin_to_cos: true, freq_shift: 0., blocks: vec![BlockConfig { out_channels: 320, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 640, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: None, attention_head_dim: 8 }], layers_per_block: 2, downsample_padding: 1, mid_block_scale_factor: 1., norm_num_groups: 32, norm_eps: 1e-5, cross_attention_dim: 1280, sliced_attention_size: None, use_linear_projection: false, time_cond_proj_dim: Some(256), time_embed_dim: None }
|
||||
}
|
||||
}
|
||||
|
||||
@@ -73,8 +79,16 @@ impl UNet2DConditionModel {
|
||||
let b_channels = config.blocks[0].out_channels;
|
||||
let bl_channels = config.blocks.last().unwrap().out_channels;
|
||||
let bl_attention_head_dim = config.blocks.last().unwrap().attention_head_dim;
|
||||
let time_embed_dim = b_channels * 4;
|
||||
let time_embed_dim = config.time_embed_dim.unwrap_or(b_channels * 4);
|
||||
let conv_cfg = nn::Conv2dConfig { padding: 1, ..Default::default() };
|
||||
|
||||
// Log important configuration details
|
||||
println!("Creating UNet2DConditionModel with {} blocks", n_blocks);
|
||||
if let Some(proj_dim) = config.time_cond_proj_dim {
|
||||
println!("UNet configured with time_cond_proj_dim: {}", proj_dim);
|
||||
println!("This means the model expects guidance embeddings of dimension {}", proj_dim);
|
||||
}
|
||||
|
||||
let conv_in = conv2d(in_channels, b_channels, 3, conv_cfg, vs.pp("conv_in"))?;
|
||||
|
||||
let time_proj = Timesteps::new(b_channels, config.flip_sin_to_cos, config.freq_shift);
|
||||
@@ -241,18 +255,138 @@ impl UNet2DConditionModel {
|
||||
|
||||
// Add guidance embedding if provided
|
||||
if let Some(guidance_emb) = guidance_embedding {
|
||||
// Print shapes for debugging
|
||||
println!("Guidance embedding shape: {:?}", guidance_emb.shape());
|
||||
println!("Time embedding shape: {:?}", time_emb.shape());
|
||||
// Get dimensions of both embeddings
|
||||
let (batch_size, time_emb_dim) = time_emb.dims2()?;
|
||||
let (guidance_batch, guidance_emb_dim) = guidance_emb.dims2()?;
|
||||
|
||||
// Get the batch size and embedding dimension
|
||||
let (batch_size, emb_dim) = time_emb.dims2()?;
|
||||
println!("==== Guidance Embedding Integration ====");
|
||||
println!("Time embedding dims: [{}, {}]", batch_size, time_emb_dim);
|
||||
println!("Guidance embedding dims: [{}, {}]", guidance_batch, guidance_emb_dim);
|
||||
|
||||
// Reshape guidance embedding to match time embedding dimensions
|
||||
let guidance_emb = guidance_emb.reshape(&[1, emb_dim])?.expand(&[batch_size, emb_dim])?;
|
||||
// Use time_cond_proj_dim from config to determine expected guidance embedding dimension
|
||||
let expected_guidance_dim = self.config.time_cond_proj_dim.unwrap_or(time_emb_dim);
|
||||
println!("Expected guidance dimension from config: {}", expected_guidance_dim);
|
||||
|
||||
// Check if dimensions match the expected projection dimension
|
||||
let processed_guidance_emb = if guidance_emb_dim != time_emb_dim {
|
||||
// Convert guidance_emb to f32 for calculations
|
||||
let guidance_emb_f32 = guidance_emb.to_dtype(DType::F32)?;
|
||||
println!("Converting guidance embedding to F32 for processing");
|
||||
|
||||
// If the guidance embedding dimension matches the expected projection dimension,
|
||||
// we need to project it to match the time embedding dimension
|
||||
if guidance_emb_dim == expected_guidance_dim {
|
||||
println!("Projecting guidance embedding from {} to {}", guidance_emb_dim, time_emb_dim);
|
||||
|
||||
// Use linear interpolation
|
||||
let src_positions: Vec<f32> = (0..guidance_emb_dim).map(|i| i as f32 / (guidance_emb_dim - 1) as f32).collect();
|
||||
let dst_positions: Vec<f32> = (0..time_emb_dim).map(|i| i as f32 / (time_emb_dim - 1) as f32).collect();
|
||||
|
||||
// Get guidance embedding data (now f32)
|
||||
let guidance_data = guidance_emb_f32.reshape((guidance_emb_dim,))?.to_vec1::<f32>()?;
|
||||
println!("Reshaping guidance data to vector of length {}", guidance_data.len());
|
||||
|
||||
// Perform linear interpolation
|
||||
let mut interpolated_data = Vec::with_capacity(time_emb_dim);
|
||||
for &t in &dst_positions {
|
||||
// Find the two source positions that surround the target position
|
||||
let mut i = 0;
|
||||
while i < src_positions.len() - 1 && src_positions[i + 1] < t {
|
||||
i += 1;
|
||||
}
|
||||
|
||||
if i == src_positions.len() - 1 {
|
||||
// At or beyond the last position
|
||||
interpolated_data.push(guidance_data[i]);
|
||||
} else if t <= src_positions[0] {
|
||||
// At or before the first position
|
||||
interpolated_data.push(guidance_data[0]);
|
||||
} else {
|
||||
// Linear interpolation between two points
|
||||
let t_norm = (t - src_positions[i]) / (src_positions[i + 1] - src_positions[i]);
|
||||
let value = guidance_data[i] * (1.0 - t_norm) + guidance_data[i + 1] * t_norm;
|
||||
interpolated_data.push(value);
|
||||
}
|
||||
}
|
||||
|
||||
println!("Created interpolated data of length {}", interpolated_data.len());
|
||||
|
||||
// Create tensor and convert to time_emb's dtype
|
||||
let tensor = Tensor::from_vec(interpolated_data, (1, time_emb_dim), device)?;
|
||||
let result = tensor.to_dtype(time_emb.dtype())?;
|
||||
println!("Final projected guidance embedding shape: {:?}", result.shape());
|
||||
result
|
||||
} else {
|
||||
// Unusual case - dimensions don't match expected pattern
|
||||
println!("Warning: Guidance embedding dimension ({}) doesn't match expected projection dimension ({})", guidance_emb_dim, expected_guidance_dim);
|
||||
|
||||
// Fallback to generic interpolation
|
||||
let guidance_data = guidance_emb_f32.reshape((guidance_emb_dim,))?.to_vec1::<f32>()?;
|
||||
|
||||
let src_positions: Vec<f32> = (0..guidance_emb_dim).map(|i| i as f32 / (guidance_emb_dim - 1) as f32).collect();
|
||||
let dst_positions: Vec<f32> = (0..time_emb_dim).map(|i| i as f32 / (time_emb_dim - 1) as f32).collect();
|
||||
|
||||
let mut interpolated_data = Vec::with_capacity(time_emb_dim);
|
||||
for &t in &dst_positions {
|
||||
let mut i = 0;
|
||||
while i < src_positions.len() - 1 && src_positions[i + 1] < t {
|
||||
i += 1;
|
||||
}
|
||||
|
||||
if i == src_positions.len() - 1 {
|
||||
interpolated_data.push(guidance_data[i]);
|
||||
} else if t <= src_positions[0] {
|
||||
interpolated_data.push(guidance_data[0]);
|
||||
} else {
|
||||
let t_norm = (t - src_positions[i]) / (src_positions[i + 1] - src_positions[i]);
|
||||
let value = guidance_data[i] * (1.0 - t_norm) + guidance_data[i + 1] * t_norm;
|
||||
interpolated_data.push(value);
|
||||
}
|
||||
}
|
||||
|
||||
println!("Created interpolated data of length {}", interpolated_data.len());
|
||||
let tensor = Tensor::from_vec(interpolated_data, (1, time_emb_dim), device)?;
|
||||
let result = tensor.to_dtype(time_emb.dtype())?;
|
||||
println!("Final fallback guidance embedding shape: {:?}", result.shape());
|
||||
result
|
||||
}
|
||||
} else {
|
||||
// Dimensions match, no processing needed
|
||||
println!("Guidance dimension already matches time embedding dimension, no interpolation needed");
|
||||
let result = guidance_emb.reshape((1, time_emb_dim))?.to_dtype(time_emb.dtype())?;
|
||||
println!("Reshaped guidance embedding shape: {:?}", result.shape());
|
||||
result
|
||||
};
|
||||
|
||||
// Expand to match batch size if needed
|
||||
let expanded_guidance = processed_guidance_emb.broadcast_as(time_emb.shape())?;
|
||||
println!("Expanded guidance embedding shape: {:?}", expanded_guidance.shape());
|
||||
|
||||
// Save the processed and expanded guidance embeddings to disk if requested
|
||||
if let Ok(save_dir) = std::env::var("CANDLE_SAVE_TENSORS_DIR") {
|
||||
// Save pre-expansion (original projection)
|
||||
let proc_path = format!("{}/unet_processed_guidance_embedding.npy", save_dir);
|
||||
println!("Saving processed guidance embedding to {}", proc_path);
|
||||
if let Err(e) = processed_guidance_emb.write_npy(&proc_path) {
|
||||
println!("Error saving processed guidance embedding: {}", e);
|
||||
} else {
|
||||
println!("Successfully saved processed guidance embedding to {}", proc_path);
|
||||
}
|
||||
|
||||
// Save post-expansion (final tensor used for addition to time_emb)
|
||||
let exp_path = format!("{}/unet_expanded_guidance_embedding.npy", save_dir);
|
||||
println!("Saving expanded guidance embedding to {}", exp_path);
|
||||
if let Err(e) = expanded_guidance.write_npy(&exp_path) {
|
||||
println!("Error saving expanded guidance embedding: {}", e);
|
||||
} else {
|
||||
println!("Successfully saved expanded guidance embedding to {}", exp_path);
|
||||
}
|
||||
}
|
||||
|
||||
// Add to time embeddings
|
||||
time_emb = (time_emb + guidance_emb)?;
|
||||
time_emb = (time_emb + expanded_guidance)?;
|
||||
println!("Final time embedding with guidance shape: {:?}", time_emb.shape());
|
||||
println!("==== End Guidance Integration ====");
|
||||
}
|
||||
|
||||
// 2. pre-process
|
||||
|
||||
@@ -4,7 +4,7 @@ import { Shape } from "konva/lib/Shape";
|
||||
import { Group } from "konva/lib/Group";
|
||||
|
||||
import { invoke } from "@tauri-apps/api/core";
|
||||
import { listen } from '@tauri-apps/api/event';
|
||||
import { listen } from "@tauri-apps/api/event";
|
||||
|
||||
import { FileUtilities } from "../FileUtilities/FileUtilities";
|
||||
import { ImageNode, VideoNode, TextNode, ShapeNode, ShapeType } from "../Nodes";
|
||||
@@ -31,7 +31,6 @@ interface ServerSettings {
|
||||
lora_path?: string;
|
||||
}
|
||||
|
||||
|
||||
interface GenerateImageParams {
|
||||
image: string;
|
||||
prompt: string;
|
||||
@@ -222,16 +221,18 @@ export class RealTimeDrawEngine {
|
||||
console.log(event);
|
||||
// Handle model download events
|
||||
const payload = event.payload as any;
|
||||
|
||||
|
||||
// Model download started
|
||||
if (payload.model_download_started) {
|
||||
const modelInfo = payload.model_download_started;
|
||||
console.log(`Model download started: ${modelInfo.model_name} (${modelInfo.model_type})`);
|
||||
|
||||
console.log(
|
||||
`Model download started: ${modelInfo.model_name} (${modelInfo.model_type})`,
|
||||
);
|
||||
|
||||
// Set up loading indicator
|
||||
isLoadingVisible.value = true;
|
||||
loadingProgress.value = 0;
|
||||
|
||||
|
||||
// Create fake progress updates
|
||||
const downloadTimer = setInterval(() => {
|
||||
loadingProgress.value += 2;
|
||||
@@ -242,23 +243,23 @@ export class RealTimeDrawEngine {
|
||||
}
|
||||
}, 500);
|
||||
}
|
||||
|
||||
|
||||
// Model download completed
|
||||
if (payload.model_download_complete) {
|
||||
const modelInfo = payload.model_download_complete;
|
||||
console.log(`Model download completed: ${modelInfo.model_name} (${modelInfo.model_type})`);
|
||||
|
||||
console.log(
|
||||
`Model download completed: ${modelInfo.model_name} (${modelInfo.model_type})`,
|
||||
);
|
||||
|
||||
// Complete the loading progress
|
||||
loadingProgress.value = 100;
|
||||
|
||||
|
||||
// Hide loading indicator after a short delay
|
||||
setTimeout(() => {
|
||||
isLoadingVisible.value = false;
|
||||
}, 1000);
|
||||
}
|
||||
|
||||
});
|
||||
|
||||
}
|
||||
|
||||
private isEnabled: boolean = false;
|
||||
@@ -332,9 +333,16 @@ export class RealTimeDrawEngine {
|
||||
const ctx = cursorCanvas.getContext("2d");
|
||||
if (!ctx) return;
|
||||
|
||||
// Draw the circle
|
||||
// Draw the outer circle with a light stroke
|
||||
ctx.beginPath();
|
||||
ctx.arc(size, size, size / 2, 0, Math.PI * 2);
|
||||
ctx.strokeStyle = "white";
|
||||
ctx.lineWidth = 3;
|
||||
ctx.stroke();
|
||||
|
||||
// Draw the inner circle with a dark stroke
|
||||
ctx.beginPath();
|
||||
ctx.arc(size, size, size / 2 - 1, 0, Math.PI * 2);
|
||||
ctx.strokeStyle = "black";
|
||||
ctx.lineWidth = 1;
|
||||
ctx.stroke();
|
||||
|
||||
@@ -3,6 +3,7 @@ import { LoadingBar, LoadingBarStatus } from "~/components/ui";
|
||||
import {
|
||||
isLoadingVisible,
|
||||
loadingProgress,
|
||||
loadingText,
|
||||
} from "~/signals/uiEvents/loadingIndicator";
|
||||
|
||||
export const SignaledLoadingIndicator = () => {
|
||||
@@ -10,6 +11,7 @@ export const SignaledLoadingIndicator = () => {
|
||||
useSignalEffect(() => {
|
||||
isLoadingVisible.value;
|
||||
loadingProgress.value;
|
||||
loadingText.value;
|
||||
});
|
||||
|
||||
if (!isLoadingVisible.value) {
|
||||
@@ -47,7 +49,7 @@ export const SignaledLoadingIndicator = () => {
|
||||
className="text-white/90"
|
||||
></path>
|
||||
</svg>
|
||||
Downloading models - {loadingProgress.value}%
|
||||
{loadingText.value} - {loadingProgress.value}%
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
|
||||
@@ -6,6 +6,21 @@
|
||||
:root {
|
||||
@apply select-none bg-ui-background font-sans text-white;
|
||||
}
|
||||
|
||||
/* Remove all focus styles */
|
||||
button {
|
||||
@apply outline-none ring-0;
|
||||
-webkit-tap-highlight-color: transparent !important;
|
||||
outline: none !important;
|
||||
box-shadow: none !important;
|
||||
}
|
||||
|
||||
button:focus {
|
||||
@apply outline-none ring-0;
|
||||
outline: none !important;
|
||||
box-shadow: none !important;
|
||||
}
|
||||
|
||||
.pegboard {
|
||||
background-image: radial-gradient(#ffffff25 1px, transparent 0);
|
||||
background-color: #1b1b1f;
|
||||
@@ -187,3 +202,30 @@
|
||||
background-position: -200% 0;
|
||||
}
|
||||
}
|
||||
|
||||
.react-colorful__saturation {
|
||||
margin: 16px 0;
|
||||
border-radius: 10px !important;
|
||||
border-bottom: none !important;
|
||||
}
|
||||
|
||||
.react-colorful__saturation-pointer {
|
||||
width: 24px !important;
|
||||
height: 24px !important;
|
||||
}
|
||||
|
||||
.react-colorful__hue,
|
||||
.react-colorful__alpha {
|
||||
height: 20px !important;
|
||||
border-radius: 50px !important;
|
||||
}
|
||||
|
||||
.react-colorful__hue {
|
||||
margin-bottom: 12px;
|
||||
}
|
||||
|
||||
.react-colorful__hue-pointer,
|
||||
.react-colorful__alpha-pointer {
|
||||
width: 20px !important;
|
||||
height: 20px !important;
|
||||
}
|
||||
@@ -58,6 +58,12 @@ export const ToolbarButton = ({
|
||||
className={mergedButtonClasses}
|
||||
disabled={disabled}
|
||||
autoFocus={false}
|
||||
tabIndex={-1}
|
||||
style={{
|
||||
outline: "none",
|
||||
boxShadow: "none",
|
||||
WebkitTapHighlightColor: "transparent",
|
||||
}}
|
||||
{...restButtonProps}
|
||||
onClick={(e) => {
|
||||
e.preventDefault();
|
||||
|
||||
@@ -28,17 +28,17 @@ export const Button = ({
|
||||
function getVariantClassNames(variant: string) {
|
||||
switch (variant) {
|
||||
case "tertiary": {
|
||||
return "bg-tertiary hover:bg-tertiary-400 text-white focus-visible:outline-tertiary-600";
|
||||
return "bg-tertiary hover:bg-tertiary-400 text-white";
|
||||
}
|
||||
case "secondary": {
|
||||
return "bg-ui-controls hover:bg-ui-controls/80 text-white focus-visible:outline-secondary";
|
||||
return "bg-ui-controls hover:bg-ui-controls/80 text-white";
|
||||
}
|
||||
case "action": {
|
||||
return " bg-action hover:bg-action-500 text-white focus-visible:outline-action";
|
||||
return "bg-action hover:bg-action-500 text-white";
|
||||
}
|
||||
case "primary":
|
||||
default: {
|
||||
return "bg-primary hover:bg-primary-600 text-white focus-visible:outline-primary-600";
|
||||
return "bg-primary hover:bg-primary-600 text-white";
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -51,7 +51,6 @@ export const Button = ({
|
||||
"text-[15px] font-semibold rounded-lg px-3 py-2 shadow-sm",
|
||||
"flex gap-2 items-center justify-center",
|
||||
"transition-all duration-150",
|
||||
"focus-visible:outline focus-visible:outline-2 focus-visible:outline-offset-2",
|
||||
getVariantClassNames(propsVariant),
|
||||
propsClassName,
|
||||
disabledClass,
|
||||
|
||||
@@ -1,12 +1,13 @@
|
||||
import { MouseEvent, MouseEventHandler, useState } from "react";
|
||||
import { HexColorPicker } from "react-colorful";
|
||||
import { EyeDropper } from "react-eyedrop";
|
||||
import { Button } from "./Button";
|
||||
import { faEyedropper } from "@fortawesome/pro-solid-svg-icons";
|
||||
import { twMerge } from "tailwind-merge";
|
||||
|
||||
type ColorPickerHTMLAttributes = Omit<React.HTMLAttributes<HTMLDivElement>, "color" | "onChange" | "onChangeCapture">;
|
||||
interface ColorPickerBaseProps<T extends string> extends ColorPickerHTMLAttributes {
|
||||
type ColorPickerHTMLAttributes = Omit<
|
||||
React.HTMLAttributes<HTMLDivElement>,
|
||||
"color" | "onChange" | "onChangeCapture"
|
||||
>;
|
||||
interface ColorPickerBaseProps<T extends string>
|
||||
extends ColorPickerHTMLAttributes {
|
||||
color: T;
|
||||
onChange: (newColor: T) => void;
|
||||
}
|
||||
@@ -17,7 +18,7 @@ export const HexEyedropPicker = ({
|
||||
dropperClassName,
|
||||
pickerClassName,
|
||||
onPickerChange,
|
||||
Picker
|
||||
Picker,
|
||||
}: {
|
||||
color: string;
|
||||
onDropperChange?: (newColor: string) => void;
|
||||
@@ -26,9 +27,8 @@ export const HexEyedropPicker = ({
|
||||
onPickerChange: (newColor: string) => void;
|
||||
Picker: (props: Partial<ColorPickerBaseProps<string>>) => JSX.Element;
|
||||
}) => {
|
||||
|
||||
const handleDropperClick = () => {
|
||||
if (!('EyeDropper' in window)) {
|
||||
if (!("EyeDropper" in window)) {
|
||||
// NOTE: This shouldn't happen cuz we're using Chromium
|
||||
console.error("EyeDropper not found in window");
|
||||
return;
|
||||
@@ -46,42 +46,40 @@ export const HexEyedropPicker = ({
|
||||
// User cancelled the picker
|
||||
return null;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
startEyeDrop().then((newColor) => {
|
||||
if (newColor) {
|
||||
onDropperChange?.(newColor);
|
||||
}
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
const [isDragging, setIsDragging] = useState(false);
|
||||
const handleOnMouseDown = () => {
|
||||
setIsDragging(true);
|
||||
}
|
||||
const handleOnMouseUp = () => {
|
||||
setIsDragging(false);
|
||||
}
|
||||
// const [isDragging, setIsDragging] = useState(false);
|
||||
// const handleOnMouseDown = () => {
|
||||
// setIsDragging(true);
|
||||
// };
|
||||
// const handleOnMouseUp = () => {
|
||||
// setIsDragging(false);
|
||||
// };
|
||||
|
||||
return (
|
||||
<div className={"relative"}>
|
||||
<Picker
|
||||
className={pickerClassName}
|
||||
color={color}
|
||||
onChange={onPickerChange}
|
||||
onMouseDown={handleOnMouseDown}
|
||||
onMouseUp={handleOnMouseUp}
|
||||
/>
|
||||
<Button
|
||||
className={twMerge([
|
||||
"rounded-full m-1 size-8 text-ui-controls bg-transparent hover:bg-black/70 hover:text-white z-20 shadow-none hover:shadow-xl duration-200",
|
||||
"z-20 size-8 rounded-xl bg-ui-background/30 text-white shadow-none duration-200 hover:bg-white/10 hover:text-white hover:shadow-xl",
|
||||
dropperClassName,
|
||||
"absolute left-0 top-0",
|
||||
isDragging ? "opacity-0 -z-10" : ""
|
||||
])}
|
||||
icon={faEyedropper}
|
||||
onClick={handleDropperClick}
|
||||
/>
|
||||
<Picker
|
||||
className={pickerClassName}
|
||||
color={color}
|
||||
onChange={onPickerChange}
|
||||
// onMouseDown={handleOnMouseDown}
|
||||
// onMouseUp={handleOnMouseUp}
|
||||
/>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
);
|
||||
};
|
||||
@@ -91,7 +91,7 @@ export const ColorPicker = ({
|
||||
if (streamChanges) {
|
||||
onChange(color);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
return (
|
||||
<Popover className="relative">
|
||||
@@ -124,7 +124,7 @@ export const ColorPicker = ({
|
||||
anchor={anchor}
|
||||
className={twMerge(
|
||||
paperWrapperStyles,
|
||||
"flex flex-col items-center gap-2 overflow-hidden",
|
||||
"flex flex-col items-center gap-2 overflow-hidden p-4",
|
||||
)}
|
||||
style={
|
||||
{
|
||||
@@ -134,13 +134,14 @@ export const ColorPicker = ({
|
||||
static
|
||||
>
|
||||
<HexEyedropPicker
|
||||
pickerClassName="overflow-hidden"
|
||||
color={currColor}
|
||||
onPickerChange={handleColorChange}
|
||||
onDropperChange={handleColorChange}
|
||||
Picker={HexAlphaColorPicker}
|
||||
/>
|
||||
<Input
|
||||
className="mt-2"
|
||||
inputClassName="bg-ui-background/30"
|
||||
style={{ width: "198px" }}
|
||||
icon={faHashtag}
|
||||
value={textInput}
|
||||
|
||||
@@ -5,13 +5,18 @@ import {
|
||||
} from "@fortawesome/pro-solid-svg-icons";
|
||||
import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
|
||||
import { Popover, PopoverButton, PopoverPanel } from "@headlessui/react";
|
||||
import { ChangeEvent, useEffect, useState } from "react";
|
||||
import { ChangeEvent, useEffect, useState, useRef } from "react";
|
||||
import { HexAlphaColorPicker } from "react-colorful";
|
||||
import { twMerge } from "tailwind-merge";
|
||||
|
||||
import { useSignals } from "@preact/signals-react/runtime";
|
||||
import { paperWrapperStyles } from "~/components/styles";
|
||||
import { BRUSH_MAX_SIZE, BRUSH_MIN_SIZE, paintBrushSize, setPaintBrushSize } from "~/signals/uiEvents/toolbarMain/paintMode";
|
||||
import {
|
||||
BRUSH_MAX_SIZE,
|
||||
BRUSH_MIN_SIZE,
|
||||
paintBrushSize,
|
||||
setPaintBrushSize,
|
||||
} from "~/signals/uiEvents/toolbarMain/paintMode";
|
||||
import { Button } from "../Button";
|
||||
import { Input } from "../Input";
|
||||
import { Slider } from "../Slider";
|
||||
@@ -52,16 +57,58 @@ export const PaintModeMenu = ({
|
||||
currColor: prevColor,
|
||||
textInput: prevColor.substring(1),
|
||||
});
|
||||
const [isMouseDown, setIsMouseDown] = useState(false);
|
||||
const [isMouseOutside, setIsMouseOutside] = useState(false);
|
||||
const closeRef = useRef<(() => void) | null>(null);
|
||||
const closeTimeoutRef = useRef<NodeJS.Timeout | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
setIsOpen(defaultOpen);
|
||||
}, [defaultOpen]);
|
||||
|
||||
// Global mouse up handler
|
||||
useEffect(() => {
|
||||
const handleGlobalMouseUp = () => {
|
||||
if (isMouseDown) {
|
||||
setIsMouseDown(false);
|
||||
|
||||
// If mouse is outside when released, close the panel after a short delay
|
||||
if (isMouseOutside && closeOnMouseLeave && closeRef.current) {
|
||||
closeTimeoutRef.current = setTimeout(() => {
|
||||
onChange(currColor);
|
||||
setIsOpen(false);
|
||||
closeRef.current?.();
|
||||
}, 500);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
window.addEventListener("mouseup", handleGlobalMouseUp);
|
||||
|
||||
return () => {
|
||||
window.removeEventListener("mouseup", handleGlobalMouseUp);
|
||||
};
|
||||
}, [isMouseDown, isMouseOutside, closeOnMouseLeave, currColor, onChange]);
|
||||
|
||||
const handleMouseLeave = (close: () => void) => {
|
||||
if (closeOnMouseLeave) {
|
||||
onChange(currColor);
|
||||
setIsOpen(false);
|
||||
close();
|
||||
closeRef.current = close;
|
||||
setIsMouseOutside(true);
|
||||
|
||||
// Only close if mouse is not down and closeOnMouseLeave is true
|
||||
if (closeOnMouseLeave && !isMouseDown) {
|
||||
closeTimeoutRef.current = setTimeout(() => {
|
||||
onChange(currColor);
|
||||
setIsOpen(false);
|
||||
close();
|
||||
}, 200);
|
||||
}
|
||||
};
|
||||
|
||||
const handleMouseEnter = () => {
|
||||
setIsMouseOutside(false);
|
||||
if (closeTimeoutRef.current) {
|
||||
clearTimeout(closeTimeoutRef.current);
|
||||
closeTimeoutRef.current = null;
|
||||
}
|
||||
};
|
||||
|
||||
@@ -92,7 +139,7 @@ export const PaintModeMenu = ({
|
||||
if (streamChanges) {
|
||||
onChange(color);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Rerender the component when these signals change:
|
||||
useSignals();
|
||||
@@ -123,12 +170,16 @@ export const PaintModeMenu = ({
|
||||
)}
|
||||
</PopoverButton>
|
||||
{(open || isOpen) && (
|
||||
<div onMouseLeave={() => handleMouseLeave(close)}>
|
||||
<div
|
||||
onMouseLeave={() => handleMouseLeave(close)}
|
||||
onMouseEnter={handleMouseEnter}
|
||||
onMouseDown={() => setIsMouseDown(true)}
|
||||
>
|
||||
<PopoverPanel
|
||||
anchor={anchor}
|
||||
className={twMerge(
|
||||
paperWrapperStyles,
|
||||
"flex flex-col items-center gap-2 overflow-hidden",
|
||||
"flex flex-col items-center gap-2 overflow-hidden p-4",
|
||||
)}
|
||||
style={
|
||||
{
|
||||
@@ -138,20 +189,23 @@ export const PaintModeMenu = ({
|
||||
static
|
||||
>
|
||||
<HexEyedropPicker
|
||||
pickerClassName="overflow-hidden"
|
||||
color={currColor}
|
||||
onPickerChange={handleColorChange}
|
||||
onDropperChange={handleColorChange}
|
||||
Picker={HexAlphaColorPicker}
|
||||
/>
|
||||
<Input
|
||||
className="mt-2"
|
||||
inputClassName="bg-ui-background/30"
|
||||
style={{ width: "198px" }}
|
||||
icon={faHashtag}
|
||||
value={textInput}
|
||||
onChange={handleTextInput}
|
||||
/>
|
||||
<div className="flex flex-col w-full mt-2 gap-2">
|
||||
<p className="text-white font-medium text-sm w-full justify-start">Brush Size:</p>
|
||||
<div className="mt-2 flex w-full flex-col gap-2">
|
||||
<p className="w-full justify-start text-sm font-medium text-white">
|
||||
Brush Size:
|
||||
</p>
|
||||
<Slider
|
||||
min={BRUSH_MIN_SIZE}
|
||||
max={BRUSH_MAX_SIZE}
|
||||
|
||||
@@ -28,13 +28,28 @@ const onLoadingProgressChanged = (callback: (progress: number) => void) => {
|
||||
});
|
||||
};
|
||||
|
||||
// Loading text signal
|
||||
export const DEFAULT_LOADING_TEXT = "Downloading";
|
||||
export const loadingText = signal<string>(DEFAULT_LOADING_TEXT);
|
||||
const setLoadingText = (text: string) => {
|
||||
loadingText.value = text;
|
||||
};
|
||||
|
||||
const onLoadingTextChanged = (callback: (text: string) => void) => {
|
||||
effect(() => {
|
||||
callback(loadingText.value);
|
||||
});
|
||||
};
|
||||
|
||||
// EXPORTS
|
||||
export const dispatchers = {
|
||||
setIsLoadingVisible,
|
||||
setLoadingProgress,
|
||||
setLoadingText,
|
||||
};
|
||||
|
||||
export const events = {
|
||||
onIsLoadingVisibleChanged,
|
||||
onLoadingProgressChanged,
|
||||
onLoadingTextChanged,
|
||||
};
|
||||
Reference in new issue
Block a user