mirror of
https://github.com/storytold/artcraft.git
synced 2026-11-05 22:04:51 -05:00
Compare commits
118
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
be184d740c | ||
|
|
de683b750a | ||
|
|
4603c65d2a | ||
|
|
c5b0d837a9 | ||
|
|
652785224a | ||
|
|
8d3ed947d1 | ||
|
|
243583f0b7 | ||
|
|
150223c48c | ||
|
|
ddf6a7f54e | ||
|
|
17c6e549b5 | ||
|
|
031d9d8f20 | ||
|
|
43df6c9181 | ||
|
|
a8e029497f | ||
|
|
4936662ddd | ||
|
|
883e32220c | ||
|
|
727eead5a7 | ||
|
|
b4cc5b84cf | ||
|
|
d7e77775ca | ||
|
|
b75472c434 | ||
|
|
a4f876daa9 | ||
|
|
4b309a5dd8 | ||
|
|
08384955e4 | ||
|
|
31037b83a3 | ||
|
|
53d93cb1ac | ||
|
|
a0de8fc603 | ||
|
|
ef3527dd85 | ||
|
|
1a887d7414 | ||
|
|
48abc0d731 | ||
|
|
0233105393 | ||
|
|
8f94f81db4 | ||
|
|
3f25aefc3f | ||
|
|
c2929cfdb6 | ||
|
|
a859bd427b | ||
|
|
706becfd7e | ||
|
|
ada1f46e62 | ||
|
|
3556dcbfd3 | ||
|
|
4ab6ca1caa | ||
|
|
74f5b14d18 | ||
|
|
ac6dcfbc05 | ||
|
|
68b0acf8ed | ||
|
|
f783910c0a | ||
|
|
f3ee145f65 | ||
|
|
f8a2f4f858 | ||
|
|
3d7e558884 | ||
|
|
51d941a35d | ||
|
|
e56ce47205 | ||
|
|
26c7b29c17 | ||
|
|
89d2898e5f | ||
|
|
8124cbbbe1 | ||
|
|
f6b3ceb537 | ||
|
|
606d1d2345 | ||
|
|
e088f8e4fd | ||
|
|
c29d6377a4 | ||
|
|
494e1e6b2b | ||
|
|
4d31f702b4 | ||
|
|
202b7876b4 | ||
|
|
360c3a268b | ||
|
|
e2d9600458 | ||
|
|
31194c6622 | ||
|
|
5045e44af7 | ||
|
|
d387a04d12 | ||
|
|
0276349b5d | ||
|
|
3c037b085b | ||
|
|
7965afd96c | ||
|
|
17dd093a23 | ||
|
|
989d36dc6c | ||
|
|
d23199c046 | ||
|
|
e67cacc23f | ||
|
|
5d9358103c | ||
|
|
c36bfeb47c | ||
|
|
ad289122e0 | ||
|
|
35a8920db8 | ||
|
|
701d84602f | ||
|
|
1d2a52556d | ||
|
|
af44de74bc | ||
|
|
920cb9e751 | ||
|
|
c9c0a1bdf7 | ||
|
|
c65b00adc8 | ||
|
|
e199086cff | ||
|
|
92d1195abf | ||
|
|
d663eed493 | ||
|
|
7c70634746 | ||
|
|
f6039c9967 | ||
|
|
8fea9c1a5e | ||
|
|
ecee0d2ea8 | ||
|
|
0d41a52d5e | ||
|
|
8ce8d111eb | ||
|
|
b1c20c7d42 | ||
|
|
000638ab9f | ||
|
|
ef8b308cf0 | ||
|
|
5ea12fd682 | ||
|
|
1d5c4df906 | ||
|
|
ae6434e11e | ||
|
|
6077232cd3 | ||
|
|
f013373180 | ||
|
|
f0c1fcf516 | ||
|
|
9f392de9c5 | ||
|
|
f5fb340de9 | ||
|
|
0392bda03c | ||
|
|
1bca478621 | ||
|
|
f9c966003a | ||
|
|
2d8245cd55 | ||
|
|
9d1d605f5c | ||
|
|
1e38fbad5c | ||
|
|
6e03df968d | ||
|
|
13e0e4bb02 | ||
|
|
442c3579e1 | ||
|
|
b8a305e271 | ||
|
|
a70061bfd3 | ||
|
|
e7a392b8e2 | ||
|
|
9147da22f2 | ||
|
|
29c2df3ee6 | ||
|
|
1db4278712 | ||
|
|
5ee910ad07 | ||
|
|
ff9d2c6b34 | ||
|
|
462dcd7645 | ||
|
|
d1617f88a4 | ||
|
|
c2e57d0815 |
@@ -66,3 +66,7 @@ scraps/
|
||||
# These should be built by the platform
|
||||
windows/sqlite/sqlite3.exp
|
||||
windows/sqlite/sqlite3.lib
|
||||
|
||||
# ArtCraft Windows DLLs
|
||||
crates/desktop/src-tauri/*.dll
|
||||
windows_dlls/
|
||||
|
||||
Generated
+3
@@ -2,6 +2,7 @@
|
||||
<dictionary name="bt">
|
||||
<words>
|
||||
<w>arpabet</w>
|
||||
<w>artcraft</w>
|
||||
<w>backfill</w>
|
||||
<w>backfilling</w>
|
||||
<w>backfills</w>
|
||||
@@ -31,6 +32,7 @@
|
||||
<w>hifigan</w>
|
||||
<w>hparams</w>
|
||||
<w>huggingface</w>
|
||||
<w>kasisnu</w>
|
||||
<w>lipsync</w>
|
||||
<w>mocap</w>
|
||||
<w>mochi</w>
|
||||
@@ -47,6 +49,7 @@
|
||||
<w>techcrunch</w>
|
||||
<w>transcoded</w>
|
||||
<w>transcodings</w>
|
||||
<w>unet</w>
|
||||
<w>unpermitted</w>
|
||||
<w>upserting</w>
|
||||
<w>upserts</w>
|
||||
|
||||
Generated
+3022
-236
File diff suppressed because it is too large
Load Diff
+11
@@ -96,6 +96,17 @@ members = [
|
||||
"crates/service/plugins/user_traits_component",
|
||||
|
||||
# Vendored code
|
||||
"crates/vendor/candle/candle-core",
|
||||
"crates/vendor/candle/candle-datasets",
|
||||
"crates/vendor/candle/candle-examples",
|
||||
"crates/vendor/candle/candle-flash-attn",
|
||||
"crates/vendor/candle/candle-kernels",
|
||||
"crates/vendor/candle/candle-metal-kernels",
|
||||
"crates/vendor/candle/candle-nn",
|
||||
"crates/vendor/candle/candle-onnx",
|
||||
"crates/vendor/candle/candle-pyo3",
|
||||
"crates/vendor/candle/candle-transformers",
|
||||
"crates/vendor/candle/tensor-tools",
|
||||
"crates/vendor/limitation",
|
||||
# NB(bt,2023-12-12): Newrelic keeps breaking with upgrades. Kill for now.
|
||||
#"crates/vendor/newrelic-telemetry-sdk-rust",
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
image_width: 1024
|
||||
image_height: 1024
|
||||
#scheduler_steps: 15
|
||||
# NB(bt): It looks like we can get down to 2 sampling steps, perhaps ?
|
||||
scheduler_steps: 8
|
||||
scheduler_samples: 1
|
||||
cfg_scale: 15.0
|
||||
seed: 42
|
||||
@@ -0,0 +1,4 @@
|
||||
# Don't restart the dev server on Rust changes
|
||||
crates/
|
||||
Cargo.toml
|
||||
Cargo.lock
|
||||
@@ -18,11 +18,73 @@ crate-type = ["staticlib", "cdylib", "rlib"]
|
||||
tauri-build = { version = "2.0.4", features = [] }
|
||||
|
||||
[dependencies]
|
||||
serde_json = "1.0"
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
log = "0.4"
|
||||
tauri = { version = "2.2.4", features = [] }
|
||||
tauri-plugin-log = "2.0.0-rc"
|
||||
once_cell = "1.20.3"
|
||||
image = "0.25.5"
|
||||
|
||||
# Internal
|
||||
filesys = { path = "../../lib/files/filesys" }
|
||||
memory_caching = { path = "../../lib/caching/memory_caching" }
|
||||
primitives = { path = "../../lib/data/primitives" }
|
||||
|
||||
# Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332
|
||||
cudarc = "0.13.7"
|
||||
|
||||
# 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"
|
||||
hf-hub = "0.3.2"
|
||||
image = "0.25.5"
|
||||
log = "0.4"
|
||||
once_cell = "1.20.3"
|
||||
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"] }
|
||||
half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
|
||||
|
||||
# NB: From `rust-background-removal`
|
||||
ndarray = "0.15.6"
|
||||
ort = { version = "1.16.2", features = ["cuda"] }
|
||||
# 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"
|
||||
|
||||
[profile.dev]
|
||||
incremental = true # Compile your binary in smaller steps.
|
||||
|
||||
[profile.release]
|
||||
codegen-units = 1 # Allows LLVM to perform better optimization.
|
||||
lto = true # Enables link-time-optimizations.
|
||||
opt-level = 3 # `s` prioritizes small binary size. Use `3` if you prefer speed.
|
||||
panic = "abort" # Higher performance by disabling panic handlers.
|
||||
strip = true # Ensures debug symbols are removed.
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
"main"
|
||||
],
|
||||
"permissions": [
|
||||
"core:default"
|
||||
"core:default",
|
||||
"log:default"
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
use crate::state::app_dir::AppDataRoot;
|
||||
use tauri::State;
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn download_models(
|
||||
_app_data_root: State<'_, AppDataRoot>,
|
||||
) -> Result<String, String> {
|
||||
|
||||
// NB: This endpoint has been removed for now as we now download models in the app lifecycle.
|
||||
|
||||
Ok(("downloaded".to_string()))
|
||||
}
|
||||
+9
-7
@@ -1,11 +1,11 @@
|
||||
use std::io::Cursor;
|
||||
use base64::prelude::BASE64_STANDARD;
|
||||
use base64::Engine;
|
||||
use base64::prelude::{BASE64_STANDARD, BASE64_URL_SAFE};
|
||||
use image::codecs::png::{CompressionType, FilterType, PngEncoder};
|
||||
use image::ImageReader;
|
||||
use std::io::Cursor;
|
||||
|
||||
#[tauri::command]
|
||||
pub fn infer_image(image: &str) -> Result<String, String> {
|
||||
pub fn flip_image(image: &str) -> Result<String, String> {
|
||||
println!("infer_image called; processing image...");
|
||||
|
||||
let bytes = BASE64_STANDARD.decode(image)
|
||||
@@ -18,17 +18,19 @@ pub fn infer_image(image: &str) -> Result<String, String> {
|
||||
.map_err(|err| format!("Image decode error: {}", err))?;
|
||||
|
||||
let rotated_image = image.rotate180();
|
||||
|
||||
|
||||
let mut buffer = Vec::new(); // TODO: Preallocate, better data structure
|
||||
let mut writer = Cursor::new(&mut buffer);
|
||||
|
||||
let encoder = PngEncoder::new_with_quality
|
||||
(&mut writer, CompressionType::Fast, FilterType::Adaptive);
|
||||
|
||||
let encoder = PngEncoder::new_with_quality(&mut writer, CompressionType::Fast, FilterType::Adaptive);
|
||||
rotated_image.write_with_encoder(encoder)
|
||||
.map_err(|err| format!("Image encoding error: {}", err))?;
|
||||
|
||||
let encoded = BASE64_STANDARD.encode(&buffer);
|
||||
|
||||
|
||||
println!("Encoded: {:?}", encoded.split_at(10).0);
|
||||
|
||||
|
||||
Ok(encoded)
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
use crate::ml::model_cache::ModelCache;
|
||||
use crate::ml::prompt_cache::PromptCache;
|
||||
use crate::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 image::imageops::FilterType;
|
||||
use image::{DynamicImage, RgbImage};
|
||||
use log::{error, info};
|
||||
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";
|
||||
|
||||
const NEGATIVE_PROMPT: &str = "bad quality, bad faces, poor quality, blurry faces, watermark";
|
||||
|
||||
const RANDOM_SEED: u32 = 42;
|
||||
|
||||
|
||||
/// This handler takes an image (as a base64 encoded string) and a prompt and returns
|
||||
/// an image (as a base64-encoded string).
|
||||
#[tauri::command]
|
||||
pub async fn infer_image(
|
||||
image: &str,
|
||||
prompt: Option<String>,
|
||||
strength: Option<f64>,
|
||||
model_config: State<'_, AppConfig>,
|
||||
model_cache: State<'_, ModelCache>,
|
||||
prompt_cache: State<'_, PromptCache>,
|
||||
app_data_root: State<'_, AppDataRoot>,
|
||||
app: AppHandle,
|
||||
) -> Result<String, String> {
|
||||
info!("infer_image endpoint called.");
|
||||
|
||||
let prompt = get_prompt_or_fallback(prompt);
|
||||
|
||||
info!("Strength: {:?}; Prompt: {}", strength, prompt);
|
||||
|
||||
let image = decode_base64_image(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 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));
|
||||
}
|
||||
};
|
||||
|
||||
let bytes = encode_rgb_image_base64_png(image)
|
||||
.map_err(|err| format!("failure to encode image: {:?}", err))?;
|
||||
|
||||
info!("Inference successful; image converted to base64, serving back to browser...");
|
||||
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
async fn do_infer_image(
|
||||
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 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,
|
||||
cfg_scale: config.cfg_scale,
|
||||
app: &app,
|
||||
app_data_root,
|
||||
use_flash_attn: true,
|
||||
};
|
||||
|
||||
let image = lcm_pipeline(args)
|
||||
.map_err(|err| format!("failure to encode image: {:?}", err))?;
|
||||
|
||||
Ok(image)
|
||||
}
|
||||
|
||||
fn get_prompt_or_fallback(user_prompt: Option<String>) -> String {
|
||||
let user_prompt = user_prompt.map(|prompt| prompt.trim().to_string())
|
||||
.filter(|prompt| !prompt.is_empty());
|
||||
|
||||
if let Some(prompt) = user_prompt {
|
||||
return prompt;
|
||||
}
|
||||
|
||||
let prompt_file = PathBuf::from(PROMPT_FILENAME)
|
||||
.canonicalize()
|
||||
.unwrap_or_else(|_| PathBuf::from(PROMPT_FILENAME));
|
||||
|
||||
std::fs::read_to_string(&prompt_file)
|
||||
.map_err(|err| format!("Failed to read prompt file: {}", err))
|
||||
.unwrap_or_else(|_| {
|
||||
error!("Failed to read prompt file: {:?}", prompt_file);
|
||||
PROMPT.to_string()
|
||||
})
|
||||
.trim()
|
||||
.to_string()
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod download_models;
|
||||
pub mod flip_image;
|
||||
pub mod image_endpoint;
|
||||
pub mod remove_background_endpoint;
|
||||
pub mod test_counter;
|
||||
@@ -0,0 +1,91 @@
|
||||
//! Adapted from https://github.com/dnanhkhoa/rust-background-removal
|
||||
|
||||
use crate::ml::background_removal::onnx_session::onnx_session;
|
||||
use crate::ml::model_cache::ModelCache;
|
||||
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};
|
||||
|
||||
/// 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))?;
|
||||
|
||||
// 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().path().join("medium.onnx"); // TODO
|
||||
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)
|
||||
})?;
|
||||
|
||||
let image = encode_dynamic_image_base64_png(image)
|
||||
.map_err(|err| format!("failure to encode image: {:?}", err))?;
|
||||
|
||||
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))
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod notification_event;
|
||||
@@ -0,0 +1,30 @@
|
||||
use serde::Serialize;
|
||||
|
||||
|
||||
#[derive(Clone, Copy, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum NotificationModelType {
|
||||
Unet,
|
||||
Vae,
|
||||
TextEncoder,
|
||||
TextTokenizer,
|
||||
}
|
||||
|
||||
#[derive(Clone, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum NotificationEvent<'a> {
|
||||
ModelDownloadStarted {
|
||||
model_name: &'a str,
|
||||
model_type: NotificationModelType,
|
||||
},
|
||||
ModelDownloadProgress {
|
||||
model_name: &'a str,
|
||||
model_type: NotificationModelType,
|
||||
currently_downloaded_bytes: usize,
|
||||
total_file_bytes: usize,
|
||||
},
|
||||
ModelDownloadComplete {
|
||||
model_name: &'a str,
|
||||
model_type: NotificationModelType,
|
||||
}
|
||||
}
|
||||
@@ -1,55 +1,72 @@
|
||||
mod image_endpoint;
|
||||
pub mod endpoints;
|
||||
pub mod events;
|
||||
pub mod ml;
|
||||
pub mod state;
|
||||
pub mod threads;
|
||||
pub mod transfer;
|
||||
pub mod utils;
|
||||
|
||||
use std::sync::{Arc, RwLock};
|
||||
use once_cell::sync::Lazy;
|
||||
use image_endpoint::infer_image;
|
||||
use crate::endpoints::download_models::download_models;
|
||||
use crate::endpoints::image_endpoint::infer_image;
|
||||
use crate::endpoints::remove_background_endpoint::remove_background;
|
||||
use crate::endpoints::test_counter::test_counter;
|
||||
use crate::ml::model_cache::ModelCache;
|
||||
use crate::ml::prompt_cache::PromptCache;
|
||||
use crate::state::app_config::AppConfig;
|
||||
use crate::threads::downloader_thread::downloader_thread;
|
||||
use tauri_plugin_log::Target;
|
||||
use tauri_plugin_log::TargetKind;
|
||||
|
||||
#[cfg_attr(mobile, tauri::mobile_entry_point)]
|
||||
pub fn run() {
|
||||
// NB: Tauri wants to install the logger itself, so we can't rely on the logger crate
|
||||
// until the tauri runtime begins.
|
||||
println!("Loading model config...");
|
||||
|
||||
let config = AppConfig::init()
|
||||
.expect("config should load");
|
||||
|
||||
let prompt_cache = PromptCache::with_capacity(8)
|
||||
.expect("prompt cache should load");
|
||||
|
||||
println!("Creating model cache...");
|
||||
|
||||
let model_cache = ModelCache::new();
|
||||
|
||||
let app_data_root = config.app_data_root.clone();
|
||||
let app_data_root2 = config.app_data_root.clone();
|
||||
|
||||
println!("Initializing backend runtime...");
|
||||
|
||||
tauri::Builder::default()
|
||||
.plugin(tauri_plugin_log::Builder::new()
|
||||
.level(log::LevelFilter::Info)
|
||||
.targets(vec![Target::new(TargetKind::Stdout)])
|
||||
.build())
|
||||
.setup(|app| {
|
||||
if cfg!(debug_assertions) {
|
||||
app.handle().plugin(
|
||||
tauri_plugin_log::Builder::default()
|
||||
.level(log::LevelFilter::Info)
|
||||
.build(),
|
||||
)?;
|
||||
}
|
||||
//if cfg!(debug_assertions) {
|
||||
// app.handle().plugin(
|
||||
// tauri_plugin_log::Builder::default()
|
||||
// .level(log::LevelFilter::Info)
|
||||
// .build(),
|
||||
// )?;
|
||||
//}
|
||||
let app = app.handle().clone();
|
||||
|
||||
tauri::async_runtime::spawn(downloader_thread(app_data_root2, app));
|
||||
|
||||
Ok(())
|
||||
})
|
||||
.manage(config)
|
||||
.manage(prompt_cache)
|
||||
.manage(model_cache)
|
||||
.manage(app_data_root)
|
||||
.invoke_handler(tauri::generate_handler![
|
||||
my_custom_command,
|
||||
test_round_trip,
|
||||
download_models,
|
||||
infer_image,
|
||||
remove_background,
|
||||
test_counter,
|
||||
])
|
||||
.run(tauri::generate_context!())
|
||||
.expect("error while running tauri application");
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn my_custom_command() {
|
||||
println!("I was invoked from JavaScript!");
|
||||
}
|
||||
|
||||
static COUNTER : Lazy<Arc<RwLock<u64>>> = Lazy::new(|| Arc::new(RwLock::new(0)));
|
||||
|
||||
#[tauri::command]
|
||||
fn test_round_trip() -> 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,6 +1,15 @@
|
||||
// Prevents additional console window on Windows in release, DO NOT REMOVE!!
|
||||
#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
|
||||
|
||||
use std::fs::File;
|
||||
|
||||
fn main() {
|
||||
//// NB: This is a test to see the behavior of the application with dynamic loading modules.
|
||||
//// We should be able to execute code before loading CUDA and other dependencies.
|
||||
//let r = File::create("program_started.txt");
|
||||
//match r {
|
||||
// Err(err) => println!("Failed to create program_started.txt: {}", err),
|
||||
// Ok(_) => {}
|
||||
//}
|
||||
app_lib::run();
|
||||
}
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
pub mod onnx_session;
|
||||
@@ -0,0 +1,18 @@
|
||||
//! Adapted from https://github.com/dnanhkhoa/rust-background-removal/blob/main/src/onnx.rs
|
||||
|
||||
use ort::{Environment, ExecutionProvider, GraphOptimizationLevel, SessionBuilder};
|
||||
|
||||
pub fn onnx_session(onnx_model_file: &str) -> 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
|
||||
.with_intra_threads(1)? // Configure the number of threads used for inference
|
||||
.with_execution_providers([
|
||||
// Configure execution providers (e.g., CUDA, CoreML, CPU)
|
||||
ExecutionProvider::CUDA(Default::default()),
|
||||
ExecutionProvider::CoreML(Default::default()),
|
||||
ExecutionProvider::CPU(Default::default()),
|
||||
])?
|
||||
.with_model_from_file(onnx_model_file)?;
|
||||
Ok(session)
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
use anyhow::Result;
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
use image::DynamicImage;
|
||||
|
||||
pub fn dynamic_image_to_tensor(image: &DynamicImage, device: &Device, datatype: DType) -> Result<Tensor> {
|
||||
let (height, width) = (image.height() as usize, image.width() as usize);
|
||||
let height = height - height % 32;
|
||||
let width = width - width % 32;
|
||||
let image = image.resize_to_fill(width as u32, height as u32, image::imageops::FilterType::CatmullRom);
|
||||
let image = image.to_rgb8();
|
||||
let image = image.into_raw();
|
||||
Ok(Tensor::from_vec(image, (height, width, 3), device)?
|
||||
.permute((2, 0, 1))?
|
||||
.to_dtype(datatype)?
|
||||
.affine(2. / 255., -1.)?
|
||||
.unsqueeze(0)?)
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
|
||||
pub fn load_image_file_to_tensor<T: AsRef<std::path::Path>>(path: T, device: &Device) -> anyhow::Result<Tensor> {
|
||||
let img = image::ImageReader::open(path)?.decode()?;
|
||||
let (height, width) = (img.height() as usize, img.width() as usize);
|
||||
println!("INPUT IMAGE WIDTH AND HEIGHT: {}x{}", width, height);
|
||||
|
||||
// TODO: Downscale from 1024 -- YES! THIS HACK DOES THE TRICK OF GETTING THE TENSOR SHAPE TO [1,4,64,64] !!
|
||||
let height = height / 2;
|
||||
let width = width / 2;
|
||||
|
||||
let height = height - height % 32;
|
||||
let width = width - width % 32;
|
||||
let img = img.resize_to_fill(
|
||||
width as u32,
|
||||
height as u32,
|
||||
image::imageops::FilterType::CatmullRom,
|
||||
);
|
||||
println!("RESIZED IMAGE WIDTH AND HEIGHT: {}x{}", width, height);
|
||||
let img = img.to_rgb8();
|
||||
let img = img.into_raw();
|
||||
let img = Tensor::from_vec(img, (height, width, 3), device)?
|
||||
.permute((2, 0, 1))?
|
||||
.to_dtype(DType::F32)?
|
||||
.affine(2. / 255., -1.)?
|
||||
.unsqueeze(0)?;
|
||||
Ok(img)
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
use anyhow::Result;
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
|
||||
pub fn load_image_file_to_tensor_2<T: AsRef<std::path::Path>>(path: T) -> Result<Tensor> {
|
||||
let img = image::ImageReader::open(path)?.decode()?;
|
||||
let (height, width) = (img.height() as usize, img.width() as usize);
|
||||
let height = height - height % 32;
|
||||
let width = width - width % 32;
|
||||
let img = img.resize_to_fill(width as u32, height as u32, image::imageops::FilterType::CatmullRom);
|
||||
let img = img.to_rgb8();
|
||||
let img = img.into_raw();
|
||||
let img = Tensor::from_vec(img, (height, width, 3), &Device::Cpu)?
|
||||
.permute((2, 0, 1))?
|
||||
.to_dtype(DType::F32)?
|
||||
.affine(2. / 255., -1.)?
|
||||
.unsqueeze(0)?;
|
||||
Ok(img)
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod dynamic_image_to_tensor;
|
||||
pub mod load_image_file_to_tensor;
|
||||
pub mod load_image_file_to_tensor_2;
|
||||
pub mod save_image_from_tensor;
|
||||
pub mod tensor_to_image_buffer;
|
||||
@@ -0,0 +1,28 @@
|
||||
use anyhow::{Error as E, Result};
|
||||
use candle_core::Tensor;
|
||||
use image::{ImageBuffer, Rgb};
|
||||
|
||||
pub fn save_image_from_tensor<P: AsRef<std::path::Path>>(img: &Tensor, p: P) -> Result<()> {
|
||||
println!("Saving image to {:?}", p.as_ref());
|
||||
let p = p.as_ref();
|
||||
let (channel, height, width) = img.dims3()?;
|
||||
println!("Image dimensions: {}x{} with {} channels", width, height, channel);
|
||||
|
||||
if channel != 3 {
|
||||
anyhow::bail!("save_image expects an input of shape (3, height, width)")
|
||||
}
|
||||
let img = img.permute((1, 2, 0))?.flatten_all()?;
|
||||
let pixels = img.to_vec1::<u8>()?;
|
||||
println!("Converting tensor to image buffer...");
|
||||
|
||||
let image: ImageBuffer<Rgb<u8>, Vec<u8>> =
|
||||
match ImageBuffer::from_raw(width as u32, height as u32, pixels) {
|
||||
Some(image) => image,
|
||||
None => anyhow::bail!("error saving image {p:?}"),
|
||||
};
|
||||
println!("Successfully created image buffer");
|
||||
|
||||
image.save(p).map_err(E::from)?;
|
||||
println!("Successfully saved image to disk");
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
use candle_core::Tensor;
|
||||
use image::{ImageBuffer, Rgb};
|
||||
|
||||
pub type RgbImage = ImageBuffer<Rgb<u8>, Vec<u8>>;
|
||||
|
||||
pub fn tensor_to_image_buffer(img: &Tensor) -> anyhow::Result<RgbImage> {
|
||||
let (channel, height, width) = img.dims3()?;
|
||||
|
||||
if channel != 3 {
|
||||
anyhow::bail!("save_image expects an input of shape (3, height, width)")
|
||||
}
|
||||
|
||||
let img = img.permute((1, 2, 0))?.flatten_all()?;
|
||||
let pixels = img.to_vec1::<u8>()?;
|
||||
|
||||
ImageBuffer::from_raw(width as u32, height as u32, pixels)
|
||||
.ok_or_else(|| anyhow::anyhow!("error loading image from raw tensor"))
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
pub mod background_removal;
|
||||
pub mod image;
|
||||
pub mod model_cache;
|
||||
pub mod model_file;
|
||||
pub mod prompt_cache;
|
||||
pub mod stable_diffusion;
|
||||
pub mod weights_registry;
|
||||
@@ -0,0 +1,60 @@
|
||||
use anyhow::anyhow;
|
||||
use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
|
||||
use candle_transformers::models::stable_diffusion::vae::AutoEncoderKL;
|
||||
use std::sync::{Arc, RwLock};
|
||||
// TODO: This data structure is gross. Generalize the locking and loading semantics for reuse.
|
||||
|
||||
// Simple registry for now. We can build complex machinery that aids in
|
||||
// VRAM utilization, disk space saving, intelligent scheduling, etc. in the future.
|
||||
pub struct ModelCache {
|
||||
vae: RwLock<Option<Arc<AutoEncoderKL>>>,
|
||||
unet2: RwLock<Option<Arc<UNet2DConditionModel>>>,
|
||||
}
|
||||
|
||||
impl ModelCache {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
vae: RwLock::new(None),
|
||||
unet2: RwLock::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_vae(&self, vae: Arc<AutoEncoderKL>) -> anyhow::Result<()> {
|
||||
match self.vae.write() {
|
||||
Err(err) => Err(anyhow!("{}", err)),
|
||||
Ok(mut lock) => {
|
||||
*lock = Some(vae);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_vae(&self) -> anyhow::Result<Option<Arc<AutoEncoderKL>>> {
|
||||
match self.vae.read() {
|
||||
Err(err) => Err(anyhow!("{}", err)),
|
||||
Ok(ref lock) => {
|
||||
Ok(lock.as_ref().map(|v|v.clone()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_unet(&self, vae: Arc<UNet2DConditionModel>) -> anyhow::Result<()> {
|
||||
match self.unet2.write() {
|
||||
Err(err) => Err(anyhow!("{}", err)),
|
||||
Ok(mut lock) => {
|
||||
*lock = Some(vae);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_unet(&self) -> anyhow::Result<Option<Arc<UNet2DConditionModel>>> {
|
||||
match self.unet2.read() {
|
||||
Err(err) => Err(anyhow!("{}", err)),
|
||||
Ok(ref lock) => {
|
||||
Ok(lock.as_ref().map(|v|v.clone()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,180 @@
|
||||
#[cfg(feature = "accelerate")]
|
||||
extern crate accelerate_src;
|
||||
#[cfg(feature = "mkl")]
|
||||
extern crate intel_mkl_src;
|
||||
|
||||
use candle_transformers::models::stable_diffusion;
|
||||
use anyhow::{Error as E, Result};
|
||||
use hf_hub::api::sync::Api;
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum StableDiffusionVersion {
|
||||
V1_5,
|
||||
V1_5Inpaint,
|
||||
V2_1,
|
||||
V2Inpaint,
|
||||
Xl,
|
||||
XlInpaint,
|
||||
Turbo,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ModelFile {
|
||||
Tokenizer,
|
||||
Tokenizer2,
|
||||
Clip,
|
||||
Clip2,
|
||||
Unet,
|
||||
UnetLcm,
|
||||
Vae,
|
||||
}
|
||||
|
||||
impl StableDiffusionVersion {
|
||||
pub fn repo(&self) -> &'static str {
|
||||
match self {
|
||||
Self::XlInpaint => "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
|
||||
Self::Xl => "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
Self::V2Inpaint => "stabilityai/stable-diffusion-2-inpainting",
|
||||
Self::V2_1 => "stabilityai/stable-diffusion-2-1",
|
||||
Self::V1_5 => "Lykon/dreamshaper-7",
|
||||
Self::V1_5Inpaint => "stable-diffusion-v1-5/stable-diffusion-inpainting",
|
||||
Self::Turbo => "Lykon/dreamshaper-7",
|
||||
}
|
||||
}
|
||||
|
||||
fn unet_file(&self, use_f16: bool) -> &'static str {
|
||||
match self {
|
||||
Self::V1_5
|
||||
| Self::V1_5Inpaint
|
||||
| Self::V2_1
|
||||
| Self::V2Inpaint
|
||||
| Self::Xl
|
||||
| Self::XlInpaint
|
||||
| Self::Turbo => {
|
||||
if use_f16 {
|
||||
"unet/diffusion_pytorch_model.fp16.safetensors"
|
||||
} else {
|
||||
"unet/diffusion_pytorch_model.safetensors"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn vae_file(&self, use_f16: bool) -> &'static str {
|
||||
match self {
|
||||
Self::V1_5
|
||||
| Self::V1_5Inpaint
|
||||
| Self::V2_1
|
||||
| Self::V2Inpaint
|
||||
| Self::Xl
|
||||
| Self::XlInpaint
|
||||
| Self::Turbo => {
|
||||
if use_f16 {
|
||||
"vae/diffusion_pytorch_model.fp16.safetensors"
|
||||
} else {
|
||||
"vae/diffusion_pytorch_model.safetensors"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn clip_file(&self, use_f16: bool) -> &'static str {
|
||||
match self {
|
||||
Self::V1_5
|
||||
| Self::V1_5Inpaint
|
||||
| Self::V2_1
|
||||
| Self::V2Inpaint
|
||||
| Self::Xl
|
||||
| Self::XlInpaint
|
||||
| Self::Turbo => {
|
||||
if use_f16 {
|
||||
"text_encoder/model.fp16.safetensors"
|
||||
} else {
|
||||
"text_encoder/model.safetensors"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn clip2_file(&self, use_f16: bool) -> &'static str {
|
||||
match self {
|
||||
Self::V1_5
|
||||
| Self::V1_5Inpaint
|
||||
| Self::V2_1
|
||||
| Self::V2Inpaint
|
||||
| Self::Xl
|
||||
| Self::XlInpaint
|
||||
| Self::Turbo => {
|
||||
if use_f16 {
|
||||
"text_encoder_2/model.fp16.safetensors"
|
||||
} else {
|
||||
"text_encoder_2/model.safetensors"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ModelFile {
|
||||
pub fn get(
|
||||
&self,
|
||||
filename: Option<String>,
|
||||
version: StableDiffusionVersion,
|
||||
use_f16: bool,
|
||||
) -> Result<std::path::PathBuf> {
|
||||
match filename {
|
||||
Some(filename) => Ok(std::path::PathBuf::from(filename)),
|
||||
None => {
|
||||
let (repo, path) = match self {
|
||||
Self::Tokenizer => {
|
||||
let tokenizer_repo = match version {
|
||||
StableDiffusionVersion::V1_5
|
||||
| StableDiffusionVersion::V2_1
|
||||
| StableDiffusionVersion::V1_5Inpaint
|
||||
| StableDiffusionVersion::V2Inpaint => "openai/clip-vit-base-patch32",
|
||||
StableDiffusionVersion::Xl
|
||||
| StableDiffusionVersion::XlInpaint
|
||||
| StableDiffusionVersion::Turbo => {
|
||||
// This seems similar to the patch32 version except some very small
|
||||
// difference in the split regex.
|
||||
"openai/clip-vit-large-patch14"
|
||||
}
|
||||
};
|
||||
(tokenizer_repo, "tokenizer.json")
|
||||
}
|
||||
Self::Tokenizer2 => {
|
||||
("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", "tokenizer.json")
|
||||
}
|
||||
Self::Clip => (version.repo(), version.clip_file(use_f16)),
|
||||
Self::Clip2 => (version.repo(), version.clip2_file(use_f16)),
|
||||
Self::UnetLcm => (
|
||||
"SimianLuo/LCM_Dreamshaper_v7",
|
||||
if use_f16 {
|
||||
"unet/diffusion_pytorch_model.fp16.safetensors"
|
||||
} else {
|
||||
"unet/diffusion_pytorch_model.safetensors"
|
||||
},
|
||||
),
|
||||
Self::Unet => (version.repo(), version.unet_file(use_f16)),
|
||||
Self::Vae => {
|
||||
// Override for SDXL when using f16 weights.
|
||||
// See https://github.com/huggingface/candle/issues/1060
|
||||
if matches!(
|
||||
version,
|
||||
StableDiffusionVersion::Xl | StableDiffusionVersion::Turbo,
|
||||
) && use_f16
|
||||
{
|
||||
(
|
||||
"madebyollin/sdxl-vae-fp16-fix",
|
||||
"diffusion_pytorch_model.safetensors",
|
||||
)
|
||||
} else {
|
||||
(version.repo(), version.vae_file(use_f16))
|
||||
}
|
||||
}
|
||||
};
|
||||
let filename = Api::new()?.model(repo.to_string()).get(path)?;
|
||||
Ok(filename)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
use candle_core::Tensor;
|
||||
use memory_caching::arc_sieve::ArcSieve;
|
||||
|
||||
pub struct PromptCache {
|
||||
cache: ArcSieve<String, Tensor>,
|
||||
}
|
||||
|
||||
impl PromptCache {
|
||||
pub fn with_capacity(capacity: usize) -> anyhow::Result<Self> {
|
||||
Ok(Self {
|
||||
cache: ArcSieve::with_capacity(capacity)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get_copy(&self, key: &String) -> anyhow::Result<Option<Tensor>> {
|
||||
self.cache.get_copy(key)
|
||||
}
|
||||
|
||||
pub fn store_copy(&self, key: &String, value: &Tensor) -> anyhow::Result<()> {
|
||||
self.cache.store_copy(key, value)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
use crate::ml::model_file::StableDiffusionVersion;
|
||||
|
||||
pub fn get_vae_scale(sd_version: StableDiffusionVersion) -> f64 {
|
||||
match sd_version {
|
||||
StableDiffusionVersion::V1_5
|
||||
| StableDiffusionVersion::V1_5Inpaint
|
||||
| StableDiffusionVersion::V2_1
|
||||
| StableDiffusionVersion::V2Inpaint
|
||||
| StableDiffusionVersion::XlInpaint
|
||||
| StableDiffusionVersion::Xl => 0.18215,
|
||||
StableDiffusionVersion::Turbo => 0.13025,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
#[cfg(feature = "accelerate")]
|
||||
extern crate accelerate_src;
|
||||
#[cfg(feature = "mkl")]
|
||||
extern crate intel_mkl_src;
|
||||
|
||||
use candle_transformers::models::stable_diffusion;
|
||||
use std::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 anyhow::Error as E;
|
||||
use candle_core::{DType, Device, Module, Tensor, D};
|
||||
use log::info;
|
||||
use tokenizers::Tokenizer;
|
||||
|
||||
#[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> {
|
||||
info!("Preparing text embeddings... for {:?}", sd_version);
|
||||
|
||||
let which = match sd_version {
|
||||
StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Turbo => vec![true, false],
|
||||
_ => vec![true],
|
||||
};
|
||||
|
||||
let text_embeddings = which
|
||||
.iter()
|
||||
.map(|first| {
|
||||
do_infer_clip_text_embeddings(
|
||||
&prompt,
|
||||
&uncond_prompt,
|
||||
None, // tokenizer
|
||||
clip_weights.clone(),
|
||||
None, // clip2_weights
|
||||
sd_version,
|
||||
&sd_config,
|
||||
false, // use_f16
|
||||
&device,
|
||||
dtype,
|
||||
use_guide_scale,
|
||||
*first,
|
||||
weights_dir,
|
||||
)
|
||||
})
|
||||
.collect::<anyhow::Result<Vec<_>>>()?;
|
||||
|
||||
let text_embeddings = Tensor::cat(&text_embeddings, D::Minus1)?;
|
||||
|
||||
Ok(text_embeddings)
|
||||
}
|
||||
|
||||
#[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> {
|
||||
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);
|
||||
|
||||
info!("Loading Clip 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(),
|
||||
};
|
||||
|
||||
println!("Running with prompt \"{prompt}\".");
|
||||
let mut tokens = tokenizer.encode(prompt, true).map_err(E::msg)?.get_ids().to_vec();
|
||||
if tokens.len() > sd_config.clip.max_position_embeddings {
|
||||
anyhow::bail!("the prompt is too long, {} > max-tokens ({})", tokens.len(), sd_config.clip.max_position_embeddings)
|
||||
}
|
||||
while tokens.len() < sd_config.clip.max_position_embeddings {
|
||||
tokens.push(pad_id)
|
||||
}
|
||||
let tokens = Tensor::new(tokens.as_slice(), device)?.unsqueeze(0)?;
|
||||
|
||||
println!("Building the Clip transformer.");
|
||||
let clip_weights = if first {
|
||||
weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER)
|
||||
} else {
|
||||
weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER_2)
|
||||
};
|
||||
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)?;
|
||||
|
||||
let text_embeddings = text_model.forward(&tokens)?;
|
||||
|
||||
let text_embeddings = if use_guide_scale {
|
||||
let mut uncond_tokens = tokenizer.encode(uncond_prompt, true).map_err(E::msg)?.get_ids().to_vec();
|
||||
if uncond_tokens.len() > sd_config.clip.max_position_embeddings {
|
||||
anyhow::bail!("the negative prompt is too long, {} > max-tokens ({})", uncond_tokens.len(), sd_config.clip.max_position_embeddings)
|
||||
}
|
||||
while uncond_tokens.len() < sd_config.clip.max_position_embeddings {
|
||||
uncond_tokens.push(pad_id)
|
||||
}
|
||||
|
||||
let uncond_tokens = Tensor::new(uncond_tokens.as_slice(), device)?.unsqueeze(0)?;
|
||||
let uncond_embeddings = text_model.forward(&uncond_tokens)?;
|
||||
|
||||
Tensor::cat(&[uncond_embeddings, text_embeddings], 0)?.to_dtype(dtype)?
|
||||
} else {
|
||||
text_embeddings.to_dtype(dtype)?
|
||||
};
|
||||
Ok(text_embeddings)
|
||||
}
|
||||
@@ -0,0 +1,258 @@
|
||||
use std::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_nn::VarBuilder;
|
||||
use candle_transformers::models::stable_diffusion::lcm::LCMScheduler;
|
||||
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 image::DynamicImage;
|
||||
use log::info;
|
||||
use rand::Rng;
|
||||
use tauri::{AppHandle, Emitter};
|
||||
|
||||
const DEFAULT_STRENGTH : f64 = 75.0;
|
||||
|
||||
pub struct Args<'a> {
|
||||
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 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 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();
|
||||
|
||||
let img2img_strength = args.i2i_strength
|
||||
.map(|s| remap_lcm_strength_range(s))
|
||||
.unwrap_or(DEFAULT_STRENGTH);
|
||||
|
||||
let img2img_strength = img2img_strength / 100.0;
|
||||
|
||||
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 seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
|
||||
configs.device.set_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;
|
||||
let use_guide_scale = false;
|
||||
|
||||
let maybe_cached = prompt_cache.get_copy(&prompt)?;
|
||||
|
||||
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,
|
||||
use_f16, // use_f16
|
||||
&configs.device,
|
||||
configs.dtype,
|
||||
use_guide_scale,
|
||||
weights_dir,
|
||||
)?;
|
||||
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 => {
|
||||
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)?;
|
||||
|
||||
let vae = Arc::new(vae);
|
||||
|
||||
model_cache.set_vae(vae.clone())?;
|
||||
|
||||
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 unet_weights = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
|
||||
|
||||
let in_channels = match configs.sd_version {
|
||||
StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::V1_5Inpaint => 9,
|
||||
_ => 4,
|
||||
};
|
||||
|
||||
// For LCM models, ensure that:
|
||||
// 1. We're using the correct prediction type (epsilon)
|
||||
// 2. The cross-attention dimension is correct (768 for Dreamshaper)
|
||||
// 3. The in_channels matches the expected input (4 for standard, 9 for inpainting)
|
||||
let unet_config = UNet2DConditionModelConfig {
|
||||
// LCM specific config values
|
||||
cross_attention_dim: 768, // Dreamshaper v7 uses 768
|
||||
time_cond_proj_dim: Some(256), // LCM models use 256-dimensional guidance embeddings
|
||||
// Rest of configuration...
|
||||
center_input_sample: false,
|
||||
flip_sin_to_cos: true,
|
||||
freq_shift: 0.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.0,
|
||||
norm_num_groups: 32,
|
||||
norm_eps: 1e-5,
|
||||
sliced_attention_size: None,
|
||||
use_linear_projection: false,
|
||||
time_embed_dim: Some(1280),
|
||||
// time_embed_dim: Some(256),
|
||||
};
|
||||
|
||||
// Load the model directly from the safetensors file
|
||||
let vs_unet = unsafe { VarBuilder::from_mmaped_safetensors(&[unet_weights], configs.dtype, &configs.device)? };
|
||||
|
||||
let unet = UNet2DConditionModel::new(vs_unet, in_channels, 4, use_flash_attn, unet_config)?;
|
||||
|
||||
let unet = Arc::new(unet);
|
||||
|
||||
model_cache.set_unet(unet.clone())?;
|
||||
|
||||
unet
|
||||
},
|
||||
};
|
||||
|
||||
let init_latent_dist: DiagonalGaussianDistribution = vae.encode(&input_image)?;
|
||||
|
||||
// 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)?;
|
||||
|
||||
// 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)?;
|
||||
|
||||
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, Some(&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)
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod get_vae_scale;
|
||||
pub mod infer_clip_text_embeddings;
|
||||
pub mod stable_diffusion_pipeline;
|
||||
pub mod lcm_pipeline;
|
||||
mod remap_lcm_strength_range;
|
||||
@@ -0,0 +1,54 @@
|
||||
use primitives::numerics::remap_range_f64::{remap_range_f64, Args};
|
||||
|
||||
/// Our LCM implementation uses strength [13, 100]. Strength below 13 will take 8 steps,
|
||||
/// so we need to limit ourselves to this range from an input range of [0, 100].
|
||||
pub fn remap_lcm_strength_range(value: f64) -> f64 {
|
||||
remap_range_f64(Args {
|
||||
old_min: 0.0,
|
||||
old_max: 100.0,
|
||||
new_min: 13.0,
|
||||
new_max: 100.0,
|
||||
value,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use crate::ml::stable_diffusion::remap_lcm_strength_range::remap_lcm_strength_range;
|
||||
use assertor::{assert_that, FloatAssertion};
|
||||
|
||||
#[test]
|
||||
fn old_max() {
|
||||
assert_eq!(remap_lcm_strength_range(100.0), 100.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn old_min() {
|
||||
assert_eq!(remap_lcm_strength_range(0.0), 13.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn under_bounds() {
|
||||
assert_eq!(remap_lcm_strength_range(-100.0), 13.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn over_bounds() {
|
||||
assert_eq!(remap_lcm_strength_range(200.0), 100.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn other_values() {
|
||||
assert_that!(remap_lcm_strength_range(0.0)).with_abs_tol(0.001).is_approx_equal_to(13.000);
|
||||
assert_that!(remap_lcm_strength_range(5.0)).with_abs_tol(0.001).is_approx_equal_to(17.350);
|
||||
assert_that!(remap_lcm_strength_range(11.0)).with_abs_tol(0.001).is_approx_equal_to(22.570);
|
||||
assert_that!(remap_lcm_strength_range(13.0)).with_abs_tol(0.001).is_approx_equal_to(24.310);
|
||||
assert_that!(remap_lcm_strength_range(14.0)).with_abs_tol(0.001).is_approx_equal_to(25.180);
|
||||
assert_that!(remap_lcm_strength_range(20.0)).with_abs_tol(0.001).is_approx_equal_to(30.400);
|
||||
assert_that!(remap_lcm_strength_range(25.0)).with_abs_tol(0.001).is_approx_equal_to(34.750);
|
||||
assert_that!(remap_lcm_strength_range(50.0)).with_abs_tol(0.001).is_approx_equal_to(56.500);
|
||||
assert_that!(remap_lcm_strength_range(75.0)).with_abs_tol(0.001).is_approx_equal_to(78.250);
|
||||
assert_that!(remap_lcm_strength_range(99.0)).with_abs_tol(0.001).is_approx_equal_to(99.130);
|
||||
assert_that!(remap_lcm_strength_range(100.0)).with_abs_tol(0.001).is_approx_equal_to(100.000);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,263 @@
|
||||
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 anyhow::{anyhow, Error as E, Result};
|
||||
use candle_core::{DType, IndexOp, Tensor, D};
|
||||
use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
|
||||
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_data_root: &'a AppDataRoot,
|
||||
pub app: &'a AppHandle,
|
||||
}
|
||||
|
||||
pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
|
||||
let Args {
|
||||
prompt,
|
||||
uncond_prompt,
|
||||
cfg_scale,
|
||||
i2i_strength,
|
||||
configs,
|
||||
model_cache,
|
||||
prompt_cache,
|
||||
app,
|
||||
image,
|
||||
app_data_root,
|
||||
} = args;
|
||||
|
||||
let weights_dir = app_data_root.weights_dir();
|
||||
|
||||
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);
|
||||
|
||||
// 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 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,
|
||||
weights_dir,
|
||||
)?;
|
||||
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 vae_file = weights_dir.weight_path(&SDXL_TURBO_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())?;
|
||||
|
||||
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 unet_weights_file = weights_dir.weight_path(&SDXL_TURBO_UNET);
|
||||
|
||||
let unet = configs.sd_config.build_unet(unet_weights_file, &configs.device, 4, true, configs.dtype)
|
||||
.map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
|
||||
|
||||
let unet = Arc::new(unet);
|
||||
|
||||
model_cache.set_unet(unet.clone())?;
|
||||
|
||||
unet
|
||||
}
|
||||
};
|
||||
|
||||
let init_latent_dist : DiagonalGaussianDistribution = vae.encode(&input_image)?;
|
||||
|
||||
// TODO(bt,2025-02-18): This takes a little bit to generate the sample.
|
||||
// This is a target for performance improvement
|
||||
println!("Generating latents from input image...");
|
||||
let latents = (init_latent_dist.sample()? * vae_scale)?;
|
||||
|
||||
//.to_device(&configs.device)?;
|
||||
|
||||
println!("Initial latents shape: {:?}", latents.shape());
|
||||
|
||||
println!("Calculating start step for diffusion process...");
|
||||
|
||||
let img2img_strength = match i2i_strength {
|
||||
None => 0.75f64,
|
||||
Some(strength) => (strength as f64) / 100.0f64,
|
||||
};
|
||||
|
||||
let t_start = {
|
||||
let start = configs.scheduler_steps - (configs.scheduler_steps as f64 * img2img_strength) as usize;
|
||||
|
||||
println!("Starting from step {} of {} (strength: {})", start, configs.scheduler_steps, img2img_strength);
|
||||
start
|
||||
};
|
||||
|
||||
let latents = if t_start < scheduler.timesteps().len() {
|
||||
println!("Adding noise to latents...");
|
||||
let noise = latents.randn_like(0f64, 1f64)?;
|
||||
scheduler.add_noise(&latents, noise, scheduler.timesteps()[t_start])?
|
||||
} else {
|
||||
println!("Using latents directly (no noise addition needed)");
|
||||
latents
|
||||
};
|
||||
|
||||
println!("Latents initialized successfully");
|
||||
|
||||
let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect();
|
||||
|
||||
// Initial noise shape: [1, 4, 64, 64]
|
||||
println!("Initial noise shape: {:?}", latents.shape());
|
||||
|
||||
println!("Starting diffusion process...");
|
||||
|
||||
let mut latents = latents;
|
||||
|
||||
for (timestep_index, ×tep) in timesteps.iter().enumerate() {
|
||||
if timestep_index < t_start {
|
||||
continue;
|
||||
}
|
||||
|
||||
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() // TODO(bt,2025-08-18): Do we have to clone the tensor? `scale_model_input` requires ownership
|
||||
};
|
||||
|
||||
let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?;
|
||||
|
||||
let mut noise_pred = match unet.forward(&latent_model_input, timestep as f64, &text_embeddings) {
|
||||
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)?)?)?;
|
||||
}
|
||||
|
||||
latents = scheduler.step(&noise_pred, timestep, &latents)?;
|
||||
}
|
||||
|
||||
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)?;
|
||||
|
||||
// TODO(bt,2025-08-18): This normalization is slow.
|
||||
// This is a target for performance improvement
|
||||
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)
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
mod weight_descriptor_builder;
|
||||
pub mod weight_descriptor;
|
||||
pub mod weights;
|
||||
@@ -0,0 +1,43 @@
|
||||
|
||||
// TODO(bt,2025-03-13): When it's time for dynamic, non-const allocation, refer to
|
||||
// http::header::name::HeaderName: https://github.com/hyperium/http/blob/master/src/header/name.rs
|
||||
// as a pattern that will support both compile-time constants and dynamically allocated String-based variants.
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub enum WeightFunction {
|
||||
TextEncoder,
|
||||
TextTokenizer,
|
||||
Unet,
|
||||
Vae,
|
||||
}
|
||||
|
||||
pub struct WeightDescriptor {
|
||||
/// Name of the weight, eg. "Lykon Dreamshaper 7"
|
||||
pub name: &'static str,
|
||||
|
||||
/// Name of the file on the filesystem, eg. lykon_dreamshaper_7_text_encoder.fp16.safetensors
|
||||
pub filename: &'static str,
|
||||
|
||||
/// URL of the file in the R2 bucket,
|
||||
/// eg. https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_text_encoder.fp16.safetensors
|
||||
pub r2_download_url: &'static str,
|
||||
|
||||
pub function: WeightFunction,
|
||||
}
|
||||
|
||||
impl WeightFunction {
|
||||
pub fn as_descriptive_str(&self) -> &'static str {
|
||||
match self {
|
||||
WeightFunction::TextEncoder => "text encoder",
|
||||
WeightFunction::TextTokenizer => "text tokenizer",
|
||||
WeightFunction::Unet => "UNET",
|
||||
WeightFunction::Vae => "VAE",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl WeightDescriptor {
|
||||
pub fn to_descriptive_name(&self) -> String {
|
||||
format!("{} ({})", self.name, self.function.as_descriptive_str())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
pub (super) const R2_BUCKET_URL : &str = "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/";
|
||||
|
||||
// TODO(bt,2025-03-13): Try to use keyword args from the call site.
|
||||
// NB: We're using a macro because we can't use a builder pattern and then concat string literals
|
||||
// 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 {
|
||||
name: $name,
|
||||
filename: $file_name,
|
||||
r2_download_url: const_format::concatcp!(crate::ml::weights_registry::weight_descriptor_builder::R2_BUCKET_URL, $file_name),
|
||||
function: $weight_function,
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub(crate) use weight;
|
||||
@@ -0,0 +1,57 @@
|
||||
use crate::ml::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction};
|
||||
use crate::ml::weights_registry::weight_descriptor_builder::weight;
|
||||
|
||||
pub const CLIP_JSON : WeightDescriptor = weight!(
|
||||
"CLIP",
|
||||
"clip_vit_bigg_14_lion2b_39b_b160k.tokenizer.json",
|
||||
WeightFunction::TextTokenizer
|
||||
);
|
||||
|
||||
/// For LCM-based image generation
|
||||
pub const LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16 : WeightDescriptor = weight!(
|
||||
"Lykon Dreamshaper 7",
|
||||
"lykon_dreamshaper_7_text_encoder.fp16.safetensors",
|
||||
WeightFunction::TextEncoder
|
||||
);
|
||||
|
||||
/// For LCM-based image generation
|
||||
pub const LYKON_DEAMSHAPER_7_VAE : WeightDescriptor = weight!(
|
||||
"Lykon Dreamshaper 7",
|
||||
"lykon_dreamshaper_7_vae.safetensors",
|
||||
WeightFunction::Vae
|
||||
);
|
||||
|
||||
/// For LCM-based image generation
|
||||
pub const SIMIANLUO_LCM_DREAMSHAPER_V7_UNET : WeightDescriptor = weight!(
|
||||
"SimianLuo LCM Dreamshaper v7",
|
||||
"simianluo_lcm_dreamshaper_v7_unet.safetensors",
|
||||
WeightFunction::Unet
|
||||
);
|
||||
|
||||
#[deprecated]
|
||||
pub const SDXL_TURBO_CLIP_TEXT_ENCODER : WeightDescriptor = weight!(
|
||||
"SDXL Turbo CLIP",
|
||||
"clip_text_encoder.safetensors",
|
||||
WeightFunction::TextEncoder
|
||||
);
|
||||
|
||||
#[deprecated]
|
||||
pub const SDXL_TURBO_CLIP_TEXT_ENCODER_2 : WeightDescriptor = weight!(
|
||||
"SDXL Turbo CLIP #2",
|
||||
"clip_text_encoder_2.safetensors",
|
||||
WeightFunction::TextEncoder
|
||||
);
|
||||
|
||||
#[deprecated]
|
||||
pub const SDXL_TURBO_VAE : WeightDescriptor = weight!(
|
||||
"SDXL Turbo",
|
||||
"diffusion_pytorch_model.vae.safetensors",
|
||||
WeightFunction::Vae
|
||||
);
|
||||
|
||||
#[deprecated]
|
||||
pub const SDXL_TURBO_UNET : WeightDescriptor = weight!(
|
||||
"SDXL Turbo",
|
||||
"diffusion_pytorch_model.unet.safetensors",
|
||||
WeightFunction::Unet
|
||||
);
|
||||
@@ -0,0 +1,109 @@
|
||||
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;
|
||||
|
||||
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,
|
||||
pub scheduler_samples: usize,
|
||||
|
||||
pub seed: Option<u64>,
|
||||
|
||||
pub cfg_scale: Option<f64>,
|
||||
|
||||
/// Probably shouldn't be here.
|
||||
pub hf_api: Api,
|
||||
|
||||
/// Root location for application data.
|
||||
pub app_data_root: AppDataRoot,
|
||||
}
|
||||
|
||||
impl AppConfig {
|
||||
|
||||
pub fn init() -> anyhow::Result<Self> {
|
||||
let yaml_configs = YamlConfig::load_from_config_file_recursive();
|
||||
|
||||
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 {
|
||||
OsPlatform::Linux => yaml_configs.linux_default_data_path.as_deref(),
|
||||
OsPlatform::MacOs => yaml_configs.mac_default_data_path.as_deref(),
|
||||
OsPlatform::Windows => yaml_configs.windows_default_data_path.as_deref(),
|
||||
};
|
||||
|
||||
println!("Possible user-defined data root: {:?}", default_data_root);
|
||||
|
||||
let app_data_root = match default_data_root {
|
||||
None => AppDataRoot::create_default()?,
|
||||
Some(path) => AppDataRoot::create_existing(path)?,
|
||||
};
|
||||
|
||||
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,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
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 std::path::{Path, PathBuf};
|
||||
use tempdir::TempDir;
|
||||
use tempfile::{Builder, NamedTempFile};
|
||||
|
||||
const DEFAULT_DATA_DIR : &str = "artcraft";
|
||||
const ASSETS_SUBDIRECTORY : &str = "assets";
|
||||
const WEIGHTS_SUBDIRECTORY : &str = "weights";
|
||||
|
||||
const TEMPORARY_SUBDIRECTORY : &str = "temp";
|
||||
|
||||
/// The path to the application data directory, which includes "asset" and "weights" data.
|
||||
#[derive(Clone)]
|
||||
pub struct AppDataRoot {
|
||||
path: PathBuf,
|
||||
assets_dir: AppAssetsDir,
|
||||
weights_dir: AppWeightsDir,
|
||||
temp_dir: TemporaryDir,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppAssetsDir {
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppWeightsDir {
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct TemporaryDir {
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
impl AppDataRoot {
|
||||
pub fn create_default() -> anyhow::Result<Self> {
|
||||
let directory = get_default_data_dir()?;
|
||||
Self::create_existing(directory)
|
||||
}
|
||||
|
||||
pub fn create_existing<P: AsRef<Path>>(dir: P) -> anyhow::Result<Self> {
|
||||
let mut dir = dir.as_ref().to_path_buf();
|
||||
|
||||
match OsPlatform::get() {
|
||||
OsPlatform::Linux | OsPlatform::MacOs => {
|
||||
if let Some(d) = dir.as_os_str().to_str() {
|
||||
dir = expanduser(d)?;
|
||||
}
|
||||
},
|
||||
OsPlatform::Windows => {}
|
||||
}
|
||||
|
||||
if !dir.is_dir() {
|
||||
println!("Creating directory {:?}", dir);
|
||||
std::fs::create_dir_all(&dir)?;
|
||||
}
|
||||
|
||||
match dir.canonicalize() {
|
||||
Ok(d) => dir = d,
|
||||
Err(err) => {
|
||||
println!("Error canonicalizing {:?}: {}", dir, err);
|
||||
}
|
||||
}
|
||||
|
||||
let assets_dir = AppAssetsDir::create_existing(dir.join(ASSETS_SUBDIRECTORY))?;
|
||||
let weights_dir = AppWeightsDir::create_existing(dir.join(WEIGHTS_SUBDIRECTORY))?;
|
||||
let temp_dir = TemporaryDir::create_existing(dir.join(TEMPORARY_SUBDIRECTORY))?;
|
||||
|
||||
Ok(Self {
|
||||
path: dir,
|
||||
assets_dir,
|
||||
weights_dir,
|
||||
temp_dir,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn assets_dir(&self) -> &AppAssetsDir {
|
||||
&self.assets_dir
|
||||
}
|
||||
|
||||
pub fn weights_dir(&self) -> &AppWeightsDir {
|
||||
&self.weights_dir
|
||||
}
|
||||
|
||||
pub fn temp_dir(&self) -> &TemporaryDir {
|
||||
&self.temp_dir
|
||||
}
|
||||
|
||||
pub fn path(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
}
|
||||
|
||||
impl AppAssetsDir {
|
||||
pub fn create_existing<P: AsRef<Path>>(dir: P) -> anyhow::Result<Self> {
|
||||
let mut dir = dir.as_ref().to_path_buf();
|
||||
match dir.canonicalize() {
|
||||
Ok(d) => dir = d,
|
||||
Err(err) => {
|
||||
println!("Error canonicalizing {:?}: {}", dir, err);
|
||||
}
|
||||
}
|
||||
if !dir.exists() {
|
||||
println!("Creating directory {:?}", dir);
|
||||
std::fs::create_dir(&dir)?;
|
||||
}
|
||||
Ok(Self {
|
||||
path: dir,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn path(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
}
|
||||
|
||||
impl AppWeightsDir {
|
||||
pub fn create_existing<P: AsRef<Path>>(dir: P) -> anyhow::Result<Self> {
|
||||
let mut dir = dir.as_ref().to_path_buf();
|
||||
match dir.canonicalize() {
|
||||
Ok(d) => dir = d,
|
||||
Err(err) => {
|
||||
println!("Error canonicalizing {:?}: {}", dir, err);
|
||||
}
|
||||
}
|
||||
if !dir.exists() {
|
||||
println!("Creating directory {:?}", dir);
|
||||
std::fs::create_dir(&dir)?;
|
||||
}
|
||||
Ok(Self {
|
||||
path: dir,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn path(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
|
||||
pub fn weight_path(&self, descriptor: &WeightDescriptor) -> PathBuf {
|
||||
self.path.join(descriptor.filename)
|
||||
}
|
||||
}
|
||||
|
||||
impl TemporaryDir {
|
||||
pub fn create_existing<P: AsRef<Path>>(dir: P) -> anyhow::Result<Self> {
|
||||
let mut dir = dir.as_ref().to_path_buf();
|
||||
match dir.canonicalize() {
|
||||
Ok(d) => dir = d,
|
||||
Err(err) => {
|
||||
println!("Error canonicalizing {:?}: {}", dir, err);
|
||||
}
|
||||
}
|
||||
if !dir.exists() {
|
||||
println!("Creating directory {:?}", dir);
|
||||
std::fs::create_dir(&dir)?;
|
||||
}
|
||||
Ok(Self {
|
||||
path: dir,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn with_prefix(&self, prefix: &str) -> anyhow::Result<TempDir> {
|
||||
let tempdir = TempDir::new_in(&self.path, prefix)?;
|
||||
Ok(tempdir)
|
||||
}
|
||||
|
||||
pub fn new_named_temp_file(&self) -> anyhow::Result<NamedTempFile> {
|
||||
let tempfile = Builder::new()
|
||||
.prefix("temp_")
|
||||
.suffix(".bin")
|
||||
.tempfile_in(&self.path)?;
|
||||
Ok(tempfile)
|
||||
}
|
||||
}
|
||||
|
||||
// eg. /home/bob/artcraft, /Users/bob/artcraft, or C:\Users\Alice\artcraft
|
||||
fn get_default_data_dir() -> anyhow::Result<PathBuf> {
|
||||
Ok(UserDirs::new()
|
||||
.ok_or_else(|| anyhow!("could not determine user home directory"))?
|
||||
.home_dir()
|
||||
.join(DEFAULT_DATA_DIR))
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
use once_cell::sync::Lazy;
|
||||
use std::io;
|
||||
use std::path::{PathBuf, MAIN_SEPARATOR};
|
||||
|
||||
static PREFIX: Lazy<String> = Lazy::new(|| format!("~{}", MAIN_SEPARATOR));
|
||||
|
||||
/// The "expanduser" crate doesn't compile on Windows, so we replace its functionality slightly
|
||||
pub fn expanduser<P: AsRef<str>>(path: P) -> io::Result<PathBuf> {
|
||||
|
||||
Ok(match path.as_ref() {
|
||||
// matches an exact "~"
|
||||
s if s == "~" => {
|
||||
home_dir()?
|
||||
},
|
||||
// matches paths that start with `~/`
|
||||
s if s.starts_with(&*PREFIX) => {
|
||||
let home = home_dir()?;
|
||||
home.join(&s[2..])
|
||||
},
|
||||
// // matches paths that start with `~` but not `~/`, might be a `~username/` path
|
||||
// s if s.starts_with("~") => {
|
||||
// let mut parts = s[1..].splitn(2, MAIN_SEPARATOR);
|
||||
// let user = parts.next()
|
||||
// .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "malformed path"))?;
|
||||
// let user = Passwd::from_name(&user)
|
||||
// .map_err(|_| io::Error::new(io::ErrorKind::Other, "error searching for user"))?
|
||||
// .ok_or_else(|| io::Error::new(io::ErrorKind::Other, format!("user '{}', does not exist", &user)))?;
|
||||
// if let Some(ref path) = parts.next() {
|
||||
// PathBuf::from(user.dir).join(&path)
|
||||
// } else {
|
||||
// PathBuf::from(user.dir)
|
||||
// }
|
||||
// },
|
||||
// nothing to expand, just make a PathBuf
|
||||
s => PathBuf::from(s)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn home_dir() -> io::Result<PathBuf> {
|
||||
dirs::home_dir()
|
||||
.ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no home directory is set"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::state::expanduser::expanduser;
|
||||
use std::env;
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[test]
|
||||
#[serial_test::serial]
|
||||
fn test_success() {
|
||||
let old_home = env::var("HOME").expect("no home dir set");
|
||||
let new_home = "/home/foo";
|
||||
env::set_var("HOME", new_home);
|
||||
let path = expanduser("~/path/to/directory");
|
||||
env::set_var("HOME", old_home);
|
||||
assert_eq!(path.expect("io error"), PathBuf::from("/home/foo/path/to/directory"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial_test::serial]
|
||||
fn test_only_tilde() {
|
||||
let old_home = env::var("HOME").expect("no home dir set");
|
||||
let new_home = "/home/foo";
|
||||
env::set_var("HOME", new_home);
|
||||
let pathstr = "~";
|
||||
let path = expanduser(pathstr);
|
||||
env::set_var("HOME", old_home);
|
||||
assert_eq!(path.expect("io error"), PathBuf::from("/home/foo"));
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod app_config;
|
||||
pub mod app_dir;
|
||||
pub mod os_platform;
|
||||
pub mod yaml_config;
|
||||
pub mod expanduser;
|
||||
@@ -0,0 +1,20 @@
|
||||
|
||||
#[derive(Copy,Clone,Eq,PartialEq,Debug)]
|
||||
pub enum OsPlatform {
|
||||
Linux,
|
||||
MacOs,
|
||||
Windows,
|
||||
}
|
||||
|
||||
impl OsPlatform {
|
||||
pub fn get() -> OsPlatform {
|
||||
// Totally scientific.
|
||||
// https://doc.rust-lang.org/std/env/consts/constant.OS.html
|
||||
match std::env::consts::OS {
|
||||
"linux" => OsPlatform::Linux,
|
||||
"windows" => OsPlatform::Windows,
|
||||
"macos" | "apple" => OsPlatform::MacOs,
|
||||
_ => OsPlatform::Linux,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
use filesys::recursively_find_file_by_name::recursively_find_file_by_name;
|
||||
use std::fs::File;
|
||||
use std::io::BufReader;
|
||||
use std::path::Path;
|
||||
|
||||
const CONFIG_FILENAME : &str = "app_config.yaml";
|
||||
|
||||
/// Configuration read or derived from a YAML config file at startup
|
||||
/// This is not meant to contain any runtime config values.
|
||||
/// This may not be exhaustive.
|
||||
#[derive(serde::Deserialize, serde::Serialize, Clone, Debug, Default)]
|
||||
pub struct YamlConfig {
|
||||
pub image_height: Option<usize>,
|
||||
pub image_width: Option<usize>,
|
||||
pub scheduler_steps: Option<usize>,
|
||||
pub scheduler_samples: Option<usize>,
|
||||
pub seed: Option<u64>,
|
||||
pub cfg_scale: Option<f64>,
|
||||
|
||||
/// Override the default location to store data on Linux
|
||||
pub linux_default_data_path: Option<String>,
|
||||
|
||||
/// Override the default location to store data on Windows
|
||||
pub windows_default_data_path: Option<String>,
|
||||
|
||||
/// Override the default location to store data on Mac
|
||||
pub mac_default_data_path: Option<String>,
|
||||
}
|
||||
|
||||
impl YamlConfig {
|
||||
pub fn read_from_file<P: AsRef<Path>>(path: P) -> anyhow::Result<Self> {
|
||||
let file = File::open(path)?;
|
||||
let reader = BufReader::new(file);
|
||||
let config = serde_yaml::from_reader(reader)?;
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
pub fn load_from_config_file_recursive() -> YamlConfig {
|
||||
let maybe_path = recursively_find_file_by_name(CONFIG_FILENAME, "../../../", 5)
|
||||
.ok()
|
||||
.flatten();
|
||||
if let Some(path) = maybe_path {
|
||||
let path = path.canonicalize().unwrap_or(path);
|
||||
println!("Attempting to load configs from: {:?}", &path);
|
||||
if let Ok(config) = YamlConfig::read_from_file(path) {
|
||||
return config;
|
||||
}
|
||||
}
|
||||
println!("Couldn't load config file. Using default configs.");
|
||||
YamlConfig::default()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
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 log::{debug, error, info};
|
||||
use std::collections::VecDeque;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::mpsc::Receiver;
|
||||
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, 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;
|
||||
const CHUNK_SIZE : usize = 1024 * 1024;
|
||||
const PARALLEL_FAILURES : usize = 8;
|
||||
const MAX_RETRIES : usize = 30;
|
||||
|
||||
struct ModelDownloadStatusChannel {
|
||||
pub model_name: String,
|
||||
pub model_type: NotificationModelType,
|
||||
pub final_file_path: PathBuf,
|
||||
pub receiver: Receiver<ProgressUpdate>,
|
||||
}
|
||||
|
||||
pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! {
|
||||
|
||||
let mut status_check_queue : Arc<Mutex<Vec<ModelDownloadStatusChannel>>> = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
{
|
||||
let status_check_queue2 = status_check_queue.clone();
|
||||
let app2 = app.clone();
|
||||
|
||||
tokio::spawn(download_notify_thread(status_check_queue2, app2));
|
||||
}
|
||||
|
||||
let mut download_queue = VecDeque::new();
|
||||
|
||||
download_queue.push_back(CLIP_JSON);
|
||||
download_queue.push_back(LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16);
|
||||
download_queue.push_back(LYKON_DEAMSHAPER_7_VAE);
|
||||
download_queue.push_back(SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
|
||||
download_queue.push_back(SDXL_TURBO_CLIP_TEXT_ENCODER); // TODO(bt): Why is this still needed?
|
||||
//download_queue.push_back(SDXL_TURBO_CLIP_TEXT_ENCODER_2); // TODO(bt): REMOVE
|
||||
//download_queue.push_back(SDXL_TURBO_VAE); // TODO(bt): REMOVE
|
||||
//download_queue.push_back(SDXL_TURBO_UNET); // TODO(bt): REMOVE
|
||||
|
||||
let mut handles = FuturesUnordered::new();
|
||||
|
||||
while let Some(model) = download_queue.pop_front() {
|
||||
let url = model.r2_download_url.to_string();
|
||||
|
||||
let filename = app_data_root.weights_dir().weight_path(&model);
|
||||
|
||||
if filename.exists() {
|
||||
info!("File already exists: {:?}", &filename);
|
||||
continue;
|
||||
}
|
||||
|
||||
let headers = None;
|
||||
|
||||
let (tx, rx) = std::sync::mpsc::channel();
|
||||
|
||||
info!("Download: {:?}", filename);
|
||||
|
||||
let temp_file = match app_data_root.temp_dir().new_named_temp_file() {
|
||||
Ok(file) => file,
|
||||
Err(err) => {
|
||||
error!("Couldn't create a temporary file: {:?}", err);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let model_type = match model.function {
|
||||
WeightFunction::TextEncoder => NotificationModelType::TextEncoder,
|
||||
WeightFunction::TextTokenizer => NotificationModelType::TextTokenizer,
|
||||
WeightFunction::Unet => NotificationModelType::Unet,
|
||||
WeightFunction::Vae => NotificationModelType::Vae,
|
||||
};
|
||||
|
||||
let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadStarted {
|
||||
model_name: model.name,
|
||||
model_type,
|
||||
});
|
||||
|
||||
if let Err(err) = result {
|
||||
error!("Could not emit event: {}", err);
|
||||
}
|
||||
|
||||
let task = tokio::spawn(download_async(
|
||||
url,
|
||||
temp_file,
|
||||
filename.clone(),
|
||||
MAX_FILES,
|
||||
CHUNK_SIZE,
|
||||
PARALLEL_FAILURES,
|
||||
MAX_RETRIES,
|
||||
headers,
|
||||
Some(tx),
|
||||
));
|
||||
|
||||
match status_check_queue.lock() {
|
||||
Err(err) => error!("Could not unlock: {:?}", err),
|
||||
Ok(mut lock) => {
|
||||
lock.push(ModelDownloadStatusChannel {
|
||||
receiver: rx,
|
||||
final_file_path: filename.clone(),
|
||||
model_name: model.to_descriptive_name(),
|
||||
model_type,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
handles.push(task);
|
||||
}
|
||||
|
||||
while let Some(result) = handles.next().await {
|
||||
match result {
|
||||
Err(err) => error!("Error downloading: {:?}", err),
|
||||
Ok(_) => {},
|
||||
}
|
||||
}
|
||||
|
||||
loop {
|
||||
// TODO(bt,2025-03-10): Queue management for new files entering the queue.
|
||||
tokio::time::sleep(std::time::Duration::from_millis(60_000)).await;
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn download_notify_thread(queue: Arc<Mutex<Vec<ModelDownloadStatusChannel>>>, app: AppHandle) -> ! {
|
||||
loop {
|
||||
match queue.lock() {
|
||||
Err(err) => error!("Could not unlock: {:?}", err),
|
||||
Ok(mut lock) => {
|
||||
debug!(">>>>> Notifier has {} items", lock.len());
|
||||
|
||||
let mut reenqueue = Vec::with_capacity(lock.len());
|
||||
|
||||
for item in (*lock).drain(..) {
|
||||
if let Ok(update) = item.receiver.try_recv() {
|
||||
let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadProgress {
|
||||
model_name: &item.model_name,
|
||||
model_type: item.model_type,
|
||||
currently_downloaded_bytes: update.complete,
|
||||
total_file_bytes: update.total_length,
|
||||
});
|
||||
|
||||
if let Err(err) = result {
|
||||
error!("Could not emit event: {}", err);
|
||||
}
|
||||
}
|
||||
|
||||
if item.final_file_path.exists() {
|
||||
debug!("No longer sending notifications for {:?}", &item.final_file_path);
|
||||
|
||||
let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadComplete {
|
||||
model_name: &item.model_name,
|
||||
model_type: item.model_type,
|
||||
});
|
||||
|
||||
if let Err(err) = result {
|
||||
error!("Could not emit event: {}", err);
|
||||
}
|
||||
} else {
|
||||
reenqueue.push(item);
|
||||
}
|
||||
}
|
||||
|
||||
lock.extend(reenqueue);
|
||||
}
|
||||
}
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod downloader_thread;
|
||||
@@ -0,0 +1,241 @@
|
||||
//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs
|
||||
|
||||
use crate::transfer::error::Error;
|
||||
use crate::transfer::util::{exponential_backoff, BASE_WAIT_TIME, MAX_WAIT_TIME};
|
||||
use futures::stream::FuturesUnordered;
|
||||
use futures::StreamExt;
|
||||
use reqwest::header::{
|
||||
HeaderMap, HeaderName, HeaderValue, AUTHORIZATION, CONTENT_RANGE,
|
||||
RANGE,
|
||||
};
|
||||
use reqwest::Url;
|
||||
use std::collections::HashMap;
|
||||
use std::fmt::Display;
|
||||
use std::fs::remove_file;
|
||||
use std::io::SeekFrom;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::mpsc::Sender;
|
||||
use std::time::Duration;
|
||||
use log::info;
|
||||
use tempfile::NamedTempFile;
|
||||
use tokio::fs::OpenOptions;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::io::AsyncSeekExt;
|
||||
use tokio::sync::Semaphore;
|
||||
use tokio::time::sleep;
|
||||
|
||||
pub struct ProgressUpdate {
|
||||
pub complete: usize,
|
||||
pub total_length: usize,
|
||||
}
|
||||
|
||||
/// max_files: Number of open file handles, which determines the maximum number of parallel downloads
|
||||
/// parallel_failures: Number of maximum failures of different chunks in parallel (cannot exceed max_files)
|
||||
/// max_retries: Number of maximum attempts per chunk. (Retries are exponentially backed off + jitter)
|
||||
///
|
||||
/// The number of threads can be tuned by the environment variable `TOKIO_WORKER_THREADS` as documented in
|
||||
/// https://docs.rs/tokio/latest/tokio/runtime/struct.Builder.html#method.worker_threads
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn download_async<P: AsRef<Path>>(
|
||||
url: String,
|
||||
filename: NamedTempFile,
|
||||
final_filename: P,
|
||||
max_files: usize,
|
||||
chunk_size: usize,
|
||||
parallel_failures: usize,
|
||||
max_retries: usize,
|
||||
input_headers: Option<HashMap<String, String>>,
|
||||
maybe_progress_sender: Option<Sender<ProgressUpdate>>,
|
||||
) -> Result<(), Error> {
|
||||
info!(">>> Downloading: {}", filename.as_ref().display());
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
// https://github.com/hyperium/hyper/issues/2136#issuecomment-589488526
|
||||
.http2_keep_alive_timeout(Duration::from_secs(15))
|
||||
.build()?;
|
||||
|
||||
let mut headers = HeaderMap::new();
|
||||
let mut auth_token = None;
|
||||
if let Some(input_headers) = input_headers {
|
||||
headers.reserve(input_headers.len());
|
||||
for (k, v) in input_headers {
|
||||
let name: HeaderName = k
|
||||
.try_into()
|
||||
.map_err(|err| Error::new_err(format!("Invalid header: {err}")))?;
|
||||
let value: HeaderValue = AsRef::<str>::as_ref(&v)
|
||||
.try_into()
|
||||
.map_err(|err| Error::new_err(format!("Invalid header value: {err}")))?;
|
||||
if name == AUTHORIZATION {
|
||||
auth_token = Some(value);
|
||||
} else {
|
||||
headers.insert(name, value);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let response = if let Some(token) = auth_token.as_ref() {
|
||||
client.get(&url).header(AUTHORIZATION, token)
|
||||
} else {
|
||||
client.get(&url)
|
||||
}
|
||||
.headers(headers.clone())
|
||||
.header(RANGE, "bytes=0-0")
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?
|
||||
.error_for_status()
|
||||
.map_err(|err| Error::new_err(err.to_string()))?;
|
||||
|
||||
// Only call the final redirect URL to avoid overloading the Hub with requests and also
|
||||
// altering the download count
|
||||
let redirected_url = response.url();
|
||||
if Url::parse(&url)
|
||||
.map_err(|err| Error::new_err(format!("failed to parse url: {err}")))?
|
||||
.host()
|
||||
== redirected_url.host()
|
||||
{
|
||||
if let Some(token) = auth_token {
|
||||
headers.insert(AUTHORIZATION, token);
|
||||
}
|
||||
}
|
||||
|
||||
let content_range = response
|
||||
.headers()
|
||||
.get(CONTENT_RANGE)
|
||||
.ok_or(Error::new_err("No content length".to_string()))?
|
||||
.to_str()
|
||||
.map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?;
|
||||
|
||||
let size: Vec<&str> = content_range.split('/').collect();
|
||||
// Content-Range: bytes 0-0/702517648
|
||||
// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Range
|
||||
let length: usize = size
|
||||
.last()
|
||||
.ok_or(Error::new_err(
|
||||
"Error while downloading: No size was detected".to_string(),
|
||||
))?
|
||||
.parse()
|
||||
.map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?;
|
||||
|
||||
|
||||
info!("File {:?} is {:?} length", filename.as_ref(), length);
|
||||
|
||||
let mut handles = FuturesUnordered::new();
|
||||
let semaphore = Arc::new(Semaphore::new(max_files));
|
||||
let parallel_failures_semaphore = Arc::new(Semaphore::new(parallel_failures));
|
||||
let max_chunk_complete = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
for start in (0..length).step_by(chunk_size) {
|
||||
let url = redirected_url.to_string();
|
||||
let filename = filename.as_ref().to_owned();
|
||||
let client = client.clone();
|
||||
let headers = headers.clone();
|
||||
|
||||
let stop = std::cmp::min(start + chunk_size - 1, length);
|
||||
let semaphore = semaphore.clone();
|
||||
let parallel_failures_semaphore = parallel_failures_semaphore.clone();
|
||||
let maybe_progress = maybe_progress_sender.clone();
|
||||
let max_chunk_complete = max_chunk_complete.clone();
|
||||
handles.push(tokio::spawn(async move {
|
||||
let permit = semaphore
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?;
|
||||
let mut chunk = download_chunk(&client, &url, &filename, start, stop, headers.clone()).await;
|
||||
let mut i = 0;
|
||||
if parallel_failures > 0 {
|
||||
while let Err(dlerr) = chunk {
|
||||
if i >= max_retries {
|
||||
return Err(Error::new_err(format!(
|
||||
"Failed after too many retries ({max_retries}): {dlerr}"
|
||||
)));
|
||||
}
|
||||
let parallel_failure_permit = parallel_failures_semaphore.clone().try_acquire_owned().map_err(|err| {
|
||||
Error::new_err(format!(
|
||||
"Failed too many failures in parallel ({parallel_failures}): {dlerr} ({err})"
|
||||
))
|
||||
})?;
|
||||
|
||||
let wait_time = exponential_backoff(BASE_WAIT_TIME, i, MAX_WAIT_TIME);
|
||||
sleep(Duration::from_millis(wait_time as u64)).await;
|
||||
|
||||
chunk = download_chunk(&client, &url, &filename, start, stop, headers.clone()).await;
|
||||
i += 1;
|
||||
drop(parallel_failure_permit);
|
||||
|
||||
if let Some(sender) = maybe_progress.as_ref() {
|
||||
// NB: Calling fetch_max twice. First call will return the lower value.
|
||||
let _old = max_chunk_complete.fetch_max(stop, Ordering::SeqCst);
|
||||
let max_chunk = max_chunk_complete.fetch_max(stop, Ordering::SeqCst);
|
||||
let progress_update = ProgressUpdate {
|
||||
complete: max_chunk,
|
||||
total_length: length,
|
||||
};
|
||||
if let Err(_err) = sender.send(progress_update) {
|
||||
// Fail silently, not essential to issue progress updates
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
drop(permit);
|
||||
chunk.map_err(|e| Error::new_err(format!("Downloading error {e}"))).and(Ok(stop - start))
|
||||
}));
|
||||
}
|
||||
|
||||
// Output the chained result
|
||||
while let Some(result) = handles.next().await {
|
||||
match result {
|
||||
Ok(Ok(size)) => {
|
||||
//if let Some(ref callback) = callback {
|
||||
// callback.call((size,), None)?;
|
||||
//}
|
||||
}
|
||||
Ok(Err(py_err)) => {
|
||||
return Err(py_err);
|
||||
}
|
||||
Err(err) => {
|
||||
return Err(Error::new_err(format!(
|
||||
"Error while downloading: {err}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let file = filename.persist(final_filename)
|
||||
.map_err(|e| Error::new_err(format!("Could not rename to final filename: {:?}", e)))?;
|
||||
|
||||
file.sync_all()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn download_chunk<P: AsRef<Path>>(
|
||||
client: &reqwest::Client,
|
||||
url: &str,
|
||||
filename: P,
|
||||
start: usize,
|
||||
stop: usize,
|
||||
headers: HeaderMap,
|
||||
) -> Result<(), Error> {
|
||||
// Process each socket concurrently.
|
||||
let range = format!("bytes={start}-{stop}");
|
||||
let mut file = OpenOptions::new()
|
||||
.write(true)
|
||||
.truncate(false)
|
||||
.create(true)
|
||||
.open(filename)
|
||||
.await?;
|
||||
file.seek(SeekFrom::Start(start as u64)).await?;
|
||||
let response = client
|
||||
.get(url)
|
||||
.headers(headers)
|
||||
.header(RANGE, range)
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
let content = response.bytes().await?;
|
||||
file.write_all(&content).await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs
|
||||
|
||||
use std::fmt::Display;
|
||||
use reqwest::header::ToStrError;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Error {
|
||||
Io(std::io::Error),
|
||||
Request(reqwest::Error),
|
||||
ToStrError(ToStrError),
|
||||
Misc(String),
|
||||
}
|
||||
|
||||
impl From<std::io::Error> for Error {
|
||||
fn from(value: std::io::Error) -> Self {
|
||||
Self::Io(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<reqwest::Error> for Error {
|
||||
fn from(value: reqwest::Error) -> Self {
|
||||
Self::Request(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ToStrError> for Error {
|
||||
fn from(value: ToStrError) -> Self {
|
||||
Self::ToStrError(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for Error {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Io(io) => write!(f, "Io: {io}"),
|
||||
Self::Request(req) => write!(f, "Request: {req}"),
|
||||
Self::ToStrError(req) => write!(f, "Response non ascii: {req}"),
|
||||
Self::Misc(err) => write!(f, "Misc error: {err}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for Error {}
|
||||
|
||||
impl Error {
|
||||
pub fn new_err(reason: String) -> Self {
|
||||
Self::Misc(reason)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
pub mod download;
|
||||
pub mod upload;
|
||||
pub mod util;
|
||||
pub mod error;
|
||||
@@ -0,0 +1,189 @@
|
||||
//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs
|
||||
|
||||
use crate::transfer::error::Error;
|
||||
use crate::transfer::util::{exponential_backoff, BASE_WAIT_TIME, MAX_WAIT_TIME};
|
||||
use futures::stream::FuturesUnordered;
|
||||
use futures::StreamExt;
|
||||
use reqwest::header::CONTENT_LENGTH;
|
||||
use std::collections::HashMap;
|
||||
use std::fmt::Display;
|
||||
use std::io::SeekFrom;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::fs::OpenOptions;
|
||||
use tokio::io::{AsyncReadExt, AsyncSeekExt};
|
||||
use tokio::sync::Semaphore;
|
||||
use tokio::time::sleep;
|
||||
use tokio_util::codec::{BytesCodec, FramedRead};
|
||||
|
||||
/// parts_urls: Dictionary consisting of part numbers as keys and the associated url as values
|
||||
/// completion_url: The url that should be called when the upload is finished
|
||||
/// max_files: Number of open file handles, which determines the maximum number of parallel uploads
|
||||
/// parallel_failures: Number of maximum failures of different chunks in parallel (cannot exceed max_files)
|
||||
/// max_retries: Number of maximum attempts per chunk. (Retries are exponentially backed off + jitter)
|
||||
///
|
||||
/// The number of threads can be tuned by the environment variable `TOKIO_WORKER_THREADS` as documented in
|
||||
/// https://docs.rs/tokio/latest/tokio/runtime/struct.Builder.html#method.worker_threads
|
||||
///
|
||||
/// See https://docs.aws.amazon.com/AmazonS3/latest/userguide/mpuoverview.html for more information
|
||||
/// on the multipart upload
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn multipart_upload(
|
||||
file_path: String,
|
||||
parts_urls: Vec<String>,
|
||||
chunk_size: u64,
|
||||
max_files: usize,
|
||||
parallel_failures: usize,
|
||||
max_retries: usize,
|
||||
) -> Result<Vec<HashMap<String, String>>, Error> {
|
||||
if parallel_failures > max_files {
|
||||
return Err(Error::new_err(
|
||||
"Error parallel_failures cannot be > max_files".to_string(),
|
||||
));
|
||||
}
|
||||
if (parallel_failures == 0) != (max_retries == 0) {
|
||||
return Err(Error::new_err(
|
||||
"For retry mechanism you need to set both `parallel_failures` and `max_retries`"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()?
|
||||
.block_on(async {
|
||||
upload_async(
|
||||
file_path,
|
||||
parts_urls,
|
||||
chunk_size,
|
||||
max_files,
|
||||
parallel_failures,
|
||||
max_retries,
|
||||
)
|
||||
.await
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn upload_async(
|
||||
file_path: String,
|
||||
parts_urls: Vec<String>,
|
||||
chunk_size: u64,
|
||||
max_files: usize,
|
||||
parallel_failures: usize,
|
||||
max_retries: usize,
|
||||
) -> Result<Vec<HashMap<String, String>>, Error> {
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let mut handles = FuturesUnordered::new();
|
||||
let semaphore = Arc::new(Semaphore::new(max_files));
|
||||
let parallel_failures_semaphore = Arc::new(Semaphore::new(parallel_failures));
|
||||
|
||||
for (part_number, part_url) in parts_urls.iter().enumerate() {
|
||||
let url = part_url.to_string();
|
||||
let path = file_path.to_owned();
|
||||
let client = client.clone();
|
||||
|
||||
let start = (part_number as u64) * chunk_size;
|
||||
let semaphore = semaphore.clone();
|
||||
let parallel_failures_semaphore = parallel_failures_semaphore.clone();
|
||||
handles.push(tokio::spawn(async move {
|
||||
let permit = semaphore
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|err| Error::new_err(format!("Error acquiring semaphore: {err}")))?;
|
||||
let mut chunk = upload_chunk(&client, &url, &path, start, chunk_size).await;
|
||||
let mut i = 0;
|
||||
if parallel_failures > 0 {
|
||||
while let Err(ul_err) = chunk {
|
||||
if i >= max_retries {
|
||||
return Err(Error::new_err(format!(
|
||||
"Failed after too many retries ({max_retries}): {ul_err}"
|
||||
)));
|
||||
}
|
||||
|
||||
let parallel_failure_permit = parallel_failures_semaphore.clone().try_acquire_owned().map_err(|err| {
|
||||
Error::new_err(format!(
|
||||
"Failed too many failures in parallel ({parallel_failures}): {ul_err} ({err})"
|
||||
))
|
||||
})?;
|
||||
|
||||
let wait_time = exponential_backoff(BASE_WAIT_TIME, i, MAX_WAIT_TIME);
|
||||
sleep(Duration::from_millis(wait_time as u64)).await;
|
||||
|
||||
chunk = upload_chunk(&client, &url, &path, start, chunk_size).await;
|
||||
i += 1;
|
||||
drop(parallel_failure_permit);
|
||||
}
|
||||
}
|
||||
drop(permit);
|
||||
chunk.map_err(|e|{
|
||||
//match e {
|
||||
// Error::Io(io) => Error::new_err(format!("Io error {io}")),
|
||||
// Error::Request(req) => Error::new_err(format!("Error while sending chunk {req}")),
|
||||
// Error::ToStrError(req) => Error::new_err(format!("Response header contains non ASCII chars: {req}")),
|
||||
//}
|
||||
e
|
||||
}).map(|chunk| (part_number, chunk, chunk_size))
|
||||
}));
|
||||
}
|
||||
|
||||
let mut results: Vec<HashMap<String, String>> = vec![HashMap::default(); parts_urls.len()];
|
||||
|
||||
while let Some(result) = handles.next().await {
|
||||
match result {
|
||||
Ok(Ok((part_number, headers, size))) => {
|
||||
//if let Some(ref callback) = callback {
|
||||
// callback.call((size,), None)?;
|
||||
//}
|
||||
results[part_number] = headers;
|
||||
}
|
||||
Ok(Err(py_err)) => {
|
||||
return Err(py_err);
|
||||
}
|
||||
Err(err) => {
|
||||
return Err(Error::new_err(format!(
|
||||
"Error occurred while uploading: {err}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
async fn upload_chunk(
|
||||
client: &reqwest::Client,
|
||||
url: &str,
|
||||
path: &str,
|
||||
start: u64,
|
||||
chunk_size: u64,
|
||||
) -> Result<HashMap<String, String>, Error> {
|
||||
let mut options = OpenOptions::new();
|
||||
let mut file = options.read(true).open(path).await?;
|
||||
let file_size = file.metadata().await?.len();
|
||||
let bytes_transferred = std::cmp::min(file_size - start, chunk_size);
|
||||
|
||||
file.seek(SeekFrom::Start(start)).await?;
|
||||
let chunk = file.take(chunk_size);
|
||||
|
||||
let response = client
|
||||
.put(url)
|
||||
.header(CONTENT_LENGTH, bytes_transferred)
|
||||
.body(reqwest::Body::wrap_stream(FramedRead::new(
|
||||
chunk,
|
||||
BytesCodec::new(),
|
||||
)))
|
||||
.send()
|
||||
.await?;
|
||||
let response = response.error_for_status()?;
|
||||
let mut headers = HashMap::new();
|
||||
for (name, value) in response.headers().into_iter() {
|
||||
headers.insert(name.to_string(), value.to_str()?.to_owned());
|
||||
}
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs
|
||||
|
||||
use rand::{thread_rng, Rng};
|
||||
use std::fmt::Display;
|
||||
|
||||
pub const BASE_WAIT_TIME: usize = 300;
|
||||
pub const MAX_WAIT_TIME: usize = 10_000;
|
||||
|
||||
fn jitter() -> usize {
|
||||
thread_rng().gen_range(0..=500)
|
||||
}
|
||||
|
||||
pub fn exponential_backoff(base_wait_time: usize, n: usize, max: usize) -> usize {
|
||||
(base_wait_time + n.pow(2) + jitter()).min(max)
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
use base64::prelude::BASE64_STANDARD;
|
||||
use base64::Engine;
|
||||
use image::{DynamicImage, ImageReader};
|
||||
use std::io::Cursor;
|
||||
|
||||
pub fn decode_base64_image(base64_image: &str) -> anyhow::Result<DynamicImage> {
|
||||
let bytes = BASE64_STANDARD.decode(base64_image)?;
|
||||
|
||||
let image = ImageReader::new(Cursor::new(bytes))
|
||||
.with_guessed_format()?
|
||||
.decode()?;
|
||||
|
||||
Ok(image)
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
use base64::prelude::BASE64_STANDARD;
|
||||
use base64::Engine;
|
||||
use image::{DynamicImage, ImageFormat};
|
||||
use std::io::Cursor;
|
||||
|
||||
pub fn encode_dynamic_image_base64_png(image: DynamicImage) -> anyhow::Result<String> {
|
||||
let mut bytes = Vec::with_capacity(1024*1024);
|
||||
|
||||
image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)?;
|
||||
|
||||
let bytes = BASE64_STANDARD.encode(bytes);
|
||||
|
||||
Ok(bytes)
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
use crate::ml::image::tensor_to_image_buffer::RgbImage;
|
||||
use base64::prelude::BASE64_STANDARD;
|
||||
use base64::Engine;
|
||||
use image::ImageFormat;
|
||||
use std::io::Cursor;
|
||||
|
||||
pub fn encode_rgb_image_base64_png(
|
||||
image: RgbImage
|
||||
) -> anyhow::Result<String> {
|
||||
let mut bytes = Vec::with_capacity(1024*1024);
|
||||
|
||||
image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)?;
|
||||
|
||||
let bytes = BASE64_STANDARD.encode(bytes);
|
||||
|
||||
Ok(bytes)
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
pub mod decode_base64_image;
|
||||
pub mod encode_dynamic_image_base64_png;
|
||||
pub mod encode_rgb_image_base64_png;
|
||||
@@ -0,0 +1 @@
|
||||
pub mod image;
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"$schema": "https://schema.tauri.app/config/2.0.0-rc",
|
||||
"productName": "Storyteller Desktop App",
|
||||
"productName": "ArtCraft",
|
||||
"version": "0.1.0",
|
||||
"identifier": "com.unique.value.thing.why",
|
||||
"build": {
|
||||
@@ -12,9 +12,9 @@
|
||||
"app": {
|
||||
"windows": [
|
||||
{
|
||||
"title": "Storyteller Desktop App Title Goes Here",
|
||||
"width": 800,
|
||||
"height": 600,
|
||||
"title": "ArtCraft",
|
||||
"width": 2400,
|
||||
"height": 1300,
|
||||
"resizable": true,
|
||||
"fullscreen": false
|
||||
}
|
||||
@@ -32,6 +32,9 @@
|
||||
"icons/128x128@2x.png",
|
||||
"icons/icon.icns",
|
||||
"icons/icon.ico"
|
||||
]
|
||||
],
|
||||
"resources": {
|
||||
"../../../windows_dlls/*.dll": "."
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
|
||||
#[test]
|
||||
fn test_dynamic_library_loading() {
|
||||
let name = libloading::library_filename("cublas");
|
||||
|
||||
assert_eq!(name, "libcublas.so");
|
||||
}
|
||||
@@ -32,3 +32,4 @@ my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" }
|
||||
|
||||
[dev-dependencies]
|
||||
testing = { path = "../../testing" }
|
||||
assertor = "0.0.3"
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
pub mod i32_to_u32_zero_clamped;
|
||||
pub mod i64_to_u64_zero_clamped;
|
||||
pub mod remap_range_f64;
|
||||
pub mod u64_to_i32_saturating;
|
||||
pub mod u64_to_i64_saturating;
|
||||
pub mod u64_to_u32_saturating;
|
||||
@@ -0,0 +1,131 @@
|
||||
|
||||
pub struct Args {
|
||||
pub old_min: f64,
|
||||
pub old_max: f64,
|
||||
pub new_min: f64,
|
||||
pub new_max: f64,
|
||||
|
||||
/// The value to remap to the new range.
|
||||
pub value: f64,
|
||||
}
|
||||
|
||||
pub fn remap_range_f64(args: Args) -> f64 {
|
||||
let old_range = args.old_max - args.old_min;
|
||||
let new_range = args.new_max - args.new_min;
|
||||
|
||||
let new_value = (args.value - args.old_min) * new_range / old_range + args.new_min;
|
||||
|
||||
if new_value > args.new_max {
|
||||
args.new_max
|
||||
} else if new_value < args.new_min {
|
||||
args.new_min
|
||||
} else {
|
||||
new_value
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::numerics::remap_range_f64::{remap_range_f64, Args};
|
||||
use assertor::{assert_that, FloatAssertion};
|
||||
|
||||
#[test]
|
||||
fn test_base_cases() {
|
||||
assert_that!(remap_range_f64(Args {
|
||||
old_min: 0.0,
|
||||
old_max: 1.0,
|
||||
new_min: 0.0,
|
||||
new_max: 1.0,
|
||||
value: 0.0,
|
||||
})).with_abs_tol(0.01).is_approx_equal_to(0.0);
|
||||
|
||||
assert_that!(remap_range_f64(Args {
|
||||
old_min: 0.0,
|
||||
old_max: 1.0,
|
||||
new_min: 0.0,
|
||||
new_max: 1.0,
|
||||
value: 1.0,
|
||||
})).with_abs_tol(0.01).is_approx_equal_to(1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_maxima() {
|
||||
assert_that!(remap_range_f64(Args {
|
||||
old_min: 0.0,
|
||||
old_max: 100.0,
|
||||
new_min: 0.0,
|
||||
new_max: 1.0,
|
||||
value: 100.0,
|
||||
})).with_abs_tol(0.01).is_approx_equal_to(1.0);
|
||||
|
||||
assert_that!(remap_range_f64(Args {
|
||||
old_min: 0.0,
|
||||
old_max: 1.0,
|
||||
new_min: 0.0,
|
||||
new_max: 50.0,
|
||||
value: 1.0,
|
||||
})).with_abs_tol(0.01).is_approx_equal_to(50.0);
|
||||
|
||||
assert_that!(remap_range_f64(Args {
|
||||
old_min: 0.0,
|
||||
old_max: 1.0,
|
||||
new_min: -100.0,
|
||||
new_max: 100.0,
|
||||
value: 1.0,
|
||||
})).with_abs_tol(0.01).is_approx_equal_to(100.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_minima() {
|
||||
assert_that!(remap_range_f64(Args {
|
||||
old_min: 0.0,
|
||||
old_max: 100.0,
|
||||
new_min: 0.0,
|
||||
new_max: 1.0,
|
||||
value: 0.0,
|
||||
})).with_abs_tol(0.01).is_approx_equal_to(0.0);
|
||||
|
||||
assert_that!(remap_range_f64(Args {
|
||||
old_min: 0.0,
|
||||
old_max: 1.0,
|
||||
new_min: 0.0,
|
||||
new_max: 50.0,
|
||||
value: 0.0,
|
||||
})).with_abs_tol(0.01).is_approx_equal_to(0.0);
|
||||
|
||||
assert_that!(remap_range_f64(Args {
|
||||
old_min: 0.0,
|
||||
old_max: 1.0,
|
||||
new_min: -100.0,
|
||||
new_max: 100.0,
|
||||
value: 0.0,
|
||||
})).with_abs_tol(0.01).is_approx_equal_to(-100.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_other_values() {
|
||||
assert_that!(remap_range_f64(Args {
|
||||
old_min: 0.0,
|
||||
old_max: 100.0,
|
||||
new_min: 0.0,
|
||||
new_max: 1.0,
|
||||
value: 50.0,
|
||||
})).with_abs_tol(0.01).is_approx_equal_to(0.50);
|
||||
|
||||
assert_that!(remap_range_f64(Args {
|
||||
old_min: 0.0,
|
||||
old_max: 1.0,
|
||||
new_min: 0.0,
|
||||
new_max: 50.0,
|
||||
value: 0.10,
|
||||
})).with_abs_tol(0.01).is_approx_equal_to(5.0);
|
||||
|
||||
assert_that!(remap_range_f64(Args {
|
||||
old_min: 0.0,
|
||||
old_max: 1.0,
|
||||
new_min: -100.0,
|
||||
new_max: 100.0,
|
||||
value: 0.25,
|
||||
})).with_abs_tol(0.01).is_approx_equal_to(-50.0);
|
||||
}
|
||||
}
|
||||
@@ -31,4 +31,5 @@ pub mod file_size;
|
||||
pub mod filename_concat;
|
||||
pub mod is_filesystem_full_error;
|
||||
pub mod path_to_string;
|
||||
pub mod recursively_find_file_by_name;
|
||||
pub mod rename_across_devices;
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
use walkdir::WalkDir;
|
||||
use errors::AnyhowResult;
|
||||
|
||||
pub fn recursively_find_file_by_name<P: AsRef<Path>, Q: AsRef<Path>>(
|
||||
search_filename: P,
|
||||
search_root: Q,
|
||||
max_depth: usize,
|
||||
) -> AnyhowResult<Option<PathBuf>> {
|
||||
|
||||
let walker = WalkDir::new(search_root)
|
||||
.max_depth(max_depth);
|
||||
|
||||
for entry in walker {
|
||||
let entry = entry?;
|
||||
if entry.path().is_dir() {
|
||||
let check_path = entry.path().join(search_filename.as_ref());
|
||||
if check_path.exists() && check_path.is_file() {
|
||||
return Ok(Some(check_path));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use testing::test_file_path::test_file_path;
|
||||
use crate::recursively_find_file_by_name::recursively_find_file_by_name;
|
||||
|
||||
#[test]
|
||||
fn test_path_exists() {
|
||||
let test_data_path = test_file_path("test_data/").expect("should work");
|
||||
let result = recursively_find_file_by_name(
|
||||
"super_mario_rpg_beware_the_forests_mushrooms.mp3",
|
||||
test_data_path,
|
||||
5
|
||||
);
|
||||
let result = result.expect("should be ok");
|
||||
let result = result.expect("should exist");
|
||||
assert!(result.ends_with("test_data/audio/mp3/super_mario_rpg_beware_the_forests_mushrooms.mp3"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_path_does_not_exist() {
|
||||
let test_data_path = test_file_path("test_data/").expect("should work");
|
||||
let result = recursively_find_file_by_name(
|
||||
"fake_name.foo",
|
||||
test_data_path,
|
||||
5
|
||||
);
|
||||
let result = result.expect("should be ok");
|
||||
assert!(result.is_none());
|
||||
}
|
||||
}
|
||||
@@ -118,24 +118,40 @@ pub static DEVELOPMENT_PREMIUM_PLANS: Lazy<HashSet<Plan>> = Lazy::new(|| {
|
||||
|
||||
/// All paid-for production premium plans in the system.
|
||||
pub static PRODUCTION_PREMIUM_PLANS: Lazy<HashSet<Plan>> = Lazy::new(|| {
|
||||
let mut plans = HashSet::new();
|
||||
let mut plans = HashSet::new();
|
||||
|
||||
// ========== All-locale plans ==========
|
||||
// ========== All-locale plans ==========
|
||||
|
||||
plans.insert(PlanBuilder::new("fakeyou_plus")
|
||||
.plan_category(PlanCategory::Paid)
|
||||
.stripe_product_id("prod_MoLv23HrxPiY7D")
|
||||
.stripe_price_id("price_1M4jDCEU5se17MekOaJ92HYX")
|
||||
.cost_per_month_dollars(7)
|
||||
.tts_base_priority_level(20)
|
||||
.tts_max_character_length(2048)
|
||||
.tts_max_duration_seconds(30)
|
||||
.web_vc_base_priority_level(20)
|
||||
.web_vc_requires_frontend_keepalive(false)
|
||||
.lipsync_requires_frontend_keepalive(false)
|
||||
.build());
|
||||
// Old "plus" price
|
||||
plans.insert(PlanBuilder::new("fakeyou_plus")
|
||||
.plan_category(PlanCategory::Paid)
|
||||
.stripe_product_id("prod_MoLv23HrxPiY7D")
|
||||
.stripe_price_id("price_1M4jDCEU5se17MekOaJ92HYX")
|
||||
.cost_per_month_dollars(7)
|
||||
.tts_base_priority_level(20)
|
||||
.tts_max_character_length(2048)
|
||||
.tts_max_duration_seconds(30)
|
||||
.web_vc_base_priority_level(20)
|
||||
.web_vc_requires_frontend_keepalive(false)
|
||||
.lipsync_requires_frontend_keepalive(false)
|
||||
.build());
|
||||
|
||||
plans.insert(PlanBuilder::new("fakeyou_pro")
|
||||
// New "plus" price
|
||||
plans.insert(PlanBuilder::new("fakeyou_plus_plan")
|
||||
.plan_category(PlanCategory::Paid)
|
||||
.stripe_product_id("prod_RvC5qUgLTSLGSD")
|
||||
.stripe_price_id("price_1R1Lh8EU5se17MekaUNtWc59")
|
||||
.cost_per_month_dollars(12)
|
||||
.tts_base_priority_level(20)
|
||||
.tts_max_character_length(2048)
|
||||
.tts_max_duration_seconds(30)
|
||||
.web_vc_base_priority_level(20)
|
||||
.web_vc_requires_frontend_keepalive(false)
|
||||
.lipsync_requires_frontend_keepalive(false)
|
||||
.build());
|
||||
|
||||
// Old "pro" price
|
||||
plans.insert(PlanBuilder::new("fakeyou_pro")
|
||||
.plan_category(PlanCategory::Paid)
|
||||
.stripe_product_id("prod_MoLw8nA6eFxHzc")
|
||||
.stripe_price_id("price_1M4jEQEU5se17MeksNfA0EKm")
|
||||
@@ -149,7 +165,23 @@ pub static PRODUCTION_PREMIUM_PLANS: Lazy<HashSet<Plan>> = Lazy::new(|| {
|
||||
.can_remove_visual_watermarks(true)
|
||||
.build());
|
||||
|
||||
plans.insert(PlanBuilder::new("fakeyou_elite")
|
||||
// New "pro" price
|
||||
plans.insert(PlanBuilder::new("fakeyou_pro_plan")
|
||||
.plan_category(PlanCategory::Paid)
|
||||
.stripe_product_id("prod_RvC6nB9rw6yAsI")
|
||||
.stripe_price_id("price_1R1LhzEU5se17Mek3JLwN3Gg")
|
||||
.cost_per_month_dollars(25)
|
||||
.tts_base_priority_level(30)
|
||||
.tts_max_character_length(3072)
|
||||
.tts_max_duration_seconds(60 * 2)
|
||||
.web_vc_base_priority_level(30)
|
||||
.web_vc_requires_frontend_keepalive(false)
|
||||
.lipsync_requires_frontend_keepalive(false)
|
||||
.can_remove_visual_watermarks(true)
|
||||
.build());
|
||||
|
||||
// Old "elite" price
|
||||
plans.insert(PlanBuilder::new("fakeyou_elite")
|
||||
.plan_category(PlanCategory::Paid)
|
||||
.stripe_product_id("prod_MoLxQmLA4R24fv")
|
||||
.stripe_price_id("price_1M4jFREU5se17Mekc8pQaSKB")
|
||||
@@ -163,6 +195,21 @@ pub static PRODUCTION_PREMIUM_PLANS: Lazy<HashSet<Plan>> = Lazy::new(|| {
|
||||
.can_remove_visual_watermarks(true)
|
||||
.build());
|
||||
|
||||
// New "elite" price
|
||||
plans.insert(PlanBuilder::new("fakeyou_elite_plan")
|
||||
.plan_category(PlanCategory::Paid)
|
||||
.stripe_product_id("prod_RvC6i84YiAYStd")
|
||||
.stripe_price_id("price_1R1LiQEU5se17MekiJdBcsHQ")
|
||||
.cost_per_month_dollars(40)
|
||||
.tts_base_priority_level(40)
|
||||
.tts_max_character_length(4096)
|
||||
.tts_max_duration_seconds(60 * 5)
|
||||
.web_vc_base_priority_level(40)
|
||||
.web_vc_requires_frontend_keepalive(false)
|
||||
.lipsync_requires_frontend_keepalive(false)
|
||||
.can_remove_visual_watermarks(true)
|
||||
.build());
|
||||
|
||||
// ========== Spanish plans ==========
|
||||
|
||||
plans.insert(PlanBuilder::new("fakeyou_basic")
|
||||
@@ -293,7 +340,7 @@ mod test {
|
||||
// making ourselves check twice when we make additions, removals, or other changes.
|
||||
#[test]
|
||||
fn test_number_of_plans_is_expected() {
|
||||
assert_eq!(11, ALL_PLANS_BY_SLUG.len());
|
||||
assert_eq!(14, ALL_PLANS_BY_SLUG.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -303,7 +350,7 @@ mod test {
|
||||
.filter(|plan| !plan.is_synthetic_plan()) // NB: Synthetic plans are production!
|
||||
.collect::<Vec<&Plan>>();
|
||||
|
||||
let expected = 4;
|
||||
let expected = 7;
|
||||
|
||||
assert_eq!(expected, production_plans.len());
|
||||
assert_eq!(expected, PRODUCTION_PREMIUM_PLANS.len());
|
||||
@@ -346,6 +393,9 @@ mod test {
|
||||
"prod_MoLw8nA6eFxHzc".to_string(), // NB: Production
|
||||
"prod_MoLxQmLA4R24fv".to_string(), // NB: Production
|
||||
"prod_MoLyt8qMDmscjr".to_string(), // NB: Production
|
||||
"prod_RvC5qUgLTSLGSD".to_string(), // NB: Production - new "plus" product
|
||||
"prod_RvC6nB9rw6yAsI".to_string(), // NB: Production - new "pro" product
|
||||
"prod_RvC6i84YiAYStd".to_string(), // NB: Production - new "elite" product
|
||||
];
|
||||
|
||||
assert_that(&product_ids).contains_all_of(&expected_product_ids.iter());
|
||||
@@ -366,6 +416,9 @@ mod test {
|
||||
"price_1M4jEQEU5se17MeksNfA0EKm".to_string(), // NB: Production
|
||||
"price_1M4jFREU5se17Mekc8pQaSKB".to_string(), // NB: Production
|
||||
"price_1M4jGPEU5se17Mek2FiztNE5".to_string(), // NB: Production
|
||||
"price_1R1Lh8EU5se17MekaUNtWc59".to_string(), // NB: Production - new "plus" price
|
||||
"price_1R1LhzEU5se17Mek3JLwN3Gg".to_string(), // NB: Production - new "pro" price
|
||||
"price_1R1LiQEU5se17MekiJdBcsHQ".to_string(), // NB: Production - new "elite" price
|
||||
];
|
||||
|
||||
assert_that(&price_ids).contains_all_of(&expected_price_ids.iter());
|
||||
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
[build]
|
||||
rustflags = ["-C", "target-cpu=native"]
|
||||
|
||||
[target.wasm32-unknown-unknown]
|
||||
rustflags = ["-C", "target-feature=+simd128"]
|
||||
|
||||
[target.x86_64-apple-darwin]
|
||||
rustflags = ["-C", "target-feature=-avx,-avx2"]
|
||||
+7
@@ -0,0 +1,7 @@
|
||||
version: 2
|
||||
updates:
|
||||
- package-ecosystem: "cargo"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
open-pull-requests-limit: 5
|
||||
@@ -0,0 +1,40 @@
|
||||
name: Deploy Rust book
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
|
||||
jobs:
|
||||
deploy:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write # To push a branch
|
||||
pull-requests: write # To create a PR from that branch
|
||||
steps:
|
||||
- uses: actions/checkout@v3
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- name: Install latest mdbook
|
||||
run: |
|
||||
tag=$(curl 'https://api.github.com/repos/rust-lang/mdbook/releases/latest' | jq -r '.tag_name')
|
||||
url="https://github.com/rust-lang/mdbook/releases/download/${tag}/mdbook-${tag}-x86_64-unknown-linux-gnu.tar.gz"
|
||||
mkdir mdbook
|
||||
curl -sSL $url | tar -xz --directory=./mdbook
|
||||
echo `pwd`/mdbook >> $GITHUB_PATH
|
||||
- name: Deploy GitHub Pages
|
||||
run: |
|
||||
# This assumes your book is in the root of your repository.
|
||||
# Just add a `cd` here if you need to change to another directory.
|
||||
cd candle-book
|
||||
mdbook build
|
||||
git worktree add gh-pages
|
||||
git config user.name "Deploy from CI"
|
||||
git config user.email ""
|
||||
cd gh-pages
|
||||
# Delete the ref to avoid keeping history.
|
||||
git update-ref -d refs/heads/gh-pages
|
||||
rm -rf *
|
||||
mv ../book/* .
|
||||
git add .
|
||||
git commit -m "Deploy $GITHUB_SHA to gh-pages"
|
||||
git push --force --set-upstream origin gh-pages
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
name: CI
|
||||
on:
|
||||
pull_request:
|
||||
|
||||
jobs:
|
||||
test:
|
||||
name: Test candle-book
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write # To push a branch
|
||||
pull-requests: write # To create a PR from that branch
|
||||
steps:
|
||||
- uses: actions/checkout@master
|
||||
- name: Install Rust
|
||||
run: |
|
||||
rustup set profile minimal
|
||||
rustup toolchain install stable
|
||||
rustup default stable
|
||||
- name: Install latest mdbook
|
||||
run: |
|
||||
tag=$(curl 'https://api.github.com/repos/rust-lang/mdbook/releases/latest' | jq -r '.tag_name')
|
||||
url="https://github.com/rust-lang/mdbook/releases/download/${tag}/mdbook-${tag}-x86_64-unknown-linux-gnu.tar.gz"
|
||||
mkdir bin
|
||||
curl -sSL $url | tar -xz --directory=bin
|
||||
echo "$(pwd)/bin" >> $GITHUB_PATH
|
||||
- name: Run tests
|
||||
run: cd candle-book && cargo build && mdbook test -L ../target/debug/deps/
|
||||
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
name: CI / cuda
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
pull_request:
|
||||
|
||||
jobs:
|
||||
test-cuda:
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.job }}-${{ github.head_ref || github.run_id }}
|
||||
cancel-in-progress: true
|
||||
runs-on:
|
||||
group: aws-g4dn-2xlarge
|
||||
container:
|
||||
image: nvidia/cuda:12.3.1-devel-ubuntu22.04
|
||||
options: --gpus 0
|
||||
if: ${{ github.event.pull_request.head.repo.full_name == github.event.pull_request.base.repo.full_name }}
|
||||
permissions:
|
||||
contents: write
|
||||
packages: write
|
||||
# This is used to complete the identity challenge
|
||||
# with sigstore/fulcio when running outside of PRs.
|
||||
id-token: write
|
||||
security-events: write
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v3
|
||||
- name: Install dependencies
|
||||
run: apt-get update && apt install curl build-essential libssl-dev protobuf-compiler pkg-config -y
|
||||
- name: Install Rust Stable
|
||||
uses: actions-rust-lang/setup-rust-toolchain@v1
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
- name: Test (cuda)
|
||||
run: cargo test --features cuda
|
||||
Binary file not shown.
@@ -0,0 +1,68 @@
|
||||
name: PyO3-CI
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths:
|
||||
- candle-pyo3/**
|
||||
pull_request:
|
||||
paths:
|
||||
- candle-pyo3/**
|
||||
|
||||
jobs:
|
||||
build_and_test:
|
||||
name: Check everything builds & tests
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
matrix:
|
||||
os: [ubuntu-latest] # For now, only test on Linux
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Install Rust
|
||||
uses: actions-rs/toolchain@v1
|
||||
with:
|
||||
toolchain: stable
|
||||
|
||||
- name: Install Python
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: 3.11
|
||||
architecture: "x64"
|
||||
|
||||
- name: Cache Cargo Registry
|
||||
uses: actions/cache@v1
|
||||
with:
|
||||
path: ~/.cargo/registry
|
||||
key: ${{ runner.os }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }}
|
||||
|
||||
- name: Install Protoc
|
||||
uses: arduino/setup-protoc@v2
|
||||
with:
|
||||
version: "25.0"
|
||||
repo-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
- name: Install
|
||||
working-directory: ./candle-pyo3
|
||||
run: |
|
||||
python -m venv .env
|
||||
source .env/bin/activate
|
||||
pip install -U pip
|
||||
pip install pytest maturin black
|
||||
python -m maturin develop -r --features onnx
|
||||
|
||||
- name: Check style
|
||||
working-directory: ./candle-pyo3
|
||||
run: |
|
||||
source .env/bin/activate
|
||||
python stub.py --check
|
||||
black --check .
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./candle-pyo3
|
||||
run: |
|
||||
source .env/bin/activate
|
||||
python -m pytest -s -v tests
|
||||
@@ -0,0 +1,87 @@
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
pull_request:
|
||||
|
||||
name: Continuous integration
|
||||
|
||||
jobs:
|
||||
check:
|
||||
name: Check
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
matrix:
|
||||
os: [ubuntu-latest, windows-latest, macOS-latest]
|
||||
rust: [stable]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
- uses: actions-rs/toolchain@v1
|
||||
with:
|
||||
profile: minimal
|
||||
toolchain: ${{ matrix.rust }}
|
||||
override: true
|
||||
- uses: actions-rs/cargo@v1
|
||||
with:
|
||||
command: check
|
||||
args: --workspace
|
||||
|
||||
test:
|
||||
name: Test Suite
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
matrix:
|
||||
os: [ubuntu-latest, windows-latest, macOS-latest]
|
||||
rust: [stable]
|
||||
steps:
|
||||
- name: Delete huge unnecessary tools folder
|
||||
if: runner.os == 'Linux'
|
||||
run: rm -rf /opt/hostedtoolcache
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
- uses: actions-rs/toolchain@v1
|
||||
with:
|
||||
profile: minimal
|
||||
toolchain: ${{ matrix.rust }}
|
||||
override: true
|
||||
- uses: actions-rs/cargo@v1
|
||||
with:
|
||||
command: test
|
||||
args: --workspace
|
||||
|
||||
fmt:
|
||||
name: Rustfmt
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions-rs/toolchain@v1
|
||||
with:
|
||||
profile: minimal
|
||||
toolchain: stable
|
||||
override: true
|
||||
- run: rustup component add rustfmt
|
||||
- uses: actions-rs/cargo@v1
|
||||
with:
|
||||
command: fmt
|
||||
args: --all -- --check
|
||||
|
||||
clippy:
|
||||
name: Clippy
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions-rs/toolchain@v1
|
||||
with:
|
||||
profile: minimal
|
||||
toolchain: stable
|
||||
override: true
|
||||
- run: rustup component add clippy
|
||||
- uses: actions-rs/cargo@v1
|
||||
with:
|
||||
command: clippy
|
||||
args: --workspace --tests --examples -- -D warnings
|
||||
@@ -0,0 +1,15 @@
|
||||
on:
|
||||
push:
|
||||
|
||||
name: Secret Leaks
|
||||
|
||||
jobs:
|
||||
trufflehog:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- name: Secret Scanning
|
||||
uses: trufflesecurity/trufflehog@main
|
||||
@@ -0,0 +1,48 @@
|
||||
# Generated by Cargo
|
||||
# will have compiled files and executables
|
||||
debug/
|
||||
data/
|
||||
dist/
|
||||
target/
|
||||
|
||||
# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
|
||||
# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
|
||||
Cargo.lock
|
||||
|
||||
# editor config
|
||||
.helix
|
||||
.vscode
|
||||
|
||||
# These are backup files generated by rustfmt
|
||||
**/*.rs.bk
|
||||
|
||||
# MSVC Windows builds of rustc generate these, which store debugging information
|
||||
*.pdb
|
||||
|
||||
*tokenizer*.json
|
||||
*.npz
|
||||
|
||||
perf.data
|
||||
flamegraph.svg
|
||||
*.dylib
|
||||
*.so
|
||||
*.swp
|
||||
*.swo
|
||||
trace-*.json
|
||||
|
||||
candle-wasm-examples/*/build
|
||||
candle-wasm-examples/*/*.bin
|
||||
candle-wasm-examples/*/*.jpeg
|
||||
candle-wasm-examples/*/audios/*.wav
|
||||
candle-wasm-examples/**/*.safetensors
|
||||
candle-wasm-examples/**/*.gguf
|
||||
candle-wasm-examples/*/package-lock.json
|
||||
candle-wasm-examples/**/config*.json
|
||||
.DS_Store
|
||||
.idea/*
|
||||
__pycache__
|
||||
out.safetensors
|
||||
out.wav
|
||||
bria.mp3
|
||||
bria.safetensors
|
||||
bria.wav
|
||||
@@ -0,0 +1,3 @@
|
||||
[submodule "candle-examples/examples/flash-attn/cutlass"]
|
||||
path = candle-flash-attn/cutlass
|
||||
url = https://github.com/NVIDIA/cutlass.git
|
||||
+15
@@ -0,0 +1,15 @@
|
||||
repos:
|
||||
- repo: https://github.com/Narsil/pre-commit-rust
|
||||
rev: 2eed6366172ef2a5186e8785ec0e67243d7d73d0
|
||||
hooks:
|
||||
- id: fmt
|
||||
name: "Rust (fmt)"
|
||||
- id: clippy
|
||||
name: "Rust (clippy)"
|
||||
args:
|
||||
[
|
||||
"--tests",
|
||||
"--examples",
|
||||
"--",
|
||||
"-Dwarnings",
|
||||
]
|
||||
Vendored
+113
@@ -0,0 +1,113 @@
|
||||
# Changelog
|
||||
This documents the main changes to the `candle` crate.
|
||||
|
||||
## v0.3.1 - Unreleased
|
||||
|
||||
### Added
|
||||
|
||||
### Modified
|
||||
|
||||
## v0.3.0 - 2023-10-01
|
||||
|
||||
### Added
|
||||
|
||||
- Added the Mistral 7b v0.1 model
|
||||
[983](https://github.com/huggingface/candle/pull/983).
|
||||
- Quantized version of the Mistral model
|
||||
[1009](https://github.com/huggingface/candle/pull/1009).
|
||||
- Add the gelu-erf op and activation function
|
||||
[969](https://github.com/huggingface/candle/pull/969).
|
||||
- Add the mixformer/phi-v1.5 model
|
||||
[930](https://github.com/huggingface/candle/pull/930).
|
||||
- Add the sclice-scatter op
|
||||
[927](https://github.com/huggingface/candle/pull/927).
|
||||
- Add the Wuerstchen diffusion model
|
||||
[911](https://github.com/huggingface/candle/pull/911).
|
||||
|
||||
### Modified
|
||||
|
||||
- Support for simd128 intrinsics in some quantized vecdots
|
||||
[982](https://github.com/huggingface/candle/pull/982).
|
||||
- Optimize the index-select cuda kernel
|
||||
[976](https://github.com/huggingface/candle/pull/976).
|
||||
- Self-contained safetensor wrappers
|
||||
[946](https://github.com/huggingface/candle/pull/946).
|
||||
|
||||
## v0.2.2 - 2023-09-18
|
||||
|
||||
### Added
|
||||
- Support for `top_p` sampling
|
||||
[819](https://github.com/huggingface/candle/pull/819).
|
||||
- T5 model including decoding
|
||||
[864](https://github.com/huggingface/candle/pull/864).
|
||||
- 1-d upsampling
|
||||
[839](https://github.com/huggingface/candle/pull/839).
|
||||
|
||||
### Modified
|
||||
- Bugfix for conv2d
|
||||
[820](https://github.com/huggingface/candle/pull/820).
|
||||
- Support tensor based indexing using `.i`
|
||||
[842](https://github.com/huggingface/candle/pull/842).
|
||||
|
||||
## v0.2.1 - 2023-09-11
|
||||
|
||||
### Added
|
||||
- Add some RNNs (GRU and LSTM) in `candle-nn`
|
||||
[674](https://github.com/huggingface/candle/pull/674),
|
||||
[688](https://github.com/huggingface/candle/pull/688).
|
||||
- gguf v2 support
|
||||
[725](https://github.com/huggingface/candle/pull/725).
|
||||
- Quantized llama example in Python using the pyo3 api
|
||||
[716](https://github.com/huggingface/candle/pull/716).
|
||||
- `candle-nn` layer for conv2d-transposed
|
||||
[760](https://github.com/huggingface/candle/pull/760).
|
||||
- Add the Segment-Anything Model (SAM) as an example
|
||||
[773](https://github.com/huggingface/candle/pull/773).
|
||||
- TinyViT backbone for the segment anything example
|
||||
[787](https://github.com/huggingface/candle/pull/787).
|
||||
- Shape with holes support
|
||||
[770](https://github.com/huggingface/candle/pull/770).
|
||||
|
||||
### Modified
|
||||
- Dilations are now supported in conv-transpose2d.
|
||||
[671](https://github.com/huggingface/candle/pull/671).
|
||||
- Interactive mode for the quantized model
|
||||
[690](https://github.com/huggingface/candle/pull/690).
|
||||
- Faster softmax operation
|
||||
[747](https://github.com/huggingface/candle/pull/747).
|
||||
- Faster convolution operations on CPU and CUDA via im2col
|
||||
[802](https://github.com/huggingface/candle/pull/802).
|
||||
- Moving some models to a more central location
|
||||
[796](https://github.com/huggingface/candle/pull/796).
|
||||
|
||||
## v0.2.0 - 2023-08-30
|
||||
|
||||
### Added
|
||||
- Add the powf op
|
||||
[664](https://github.com/huggingface/candle/pull/664).
|
||||
- Stable Diffusion XL support
|
||||
[647](https://github.com/huggingface/candle/pull/647).
|
||||
- Add the conv-transpose2d op
|
||||
[635](https://github.com/huggingface/candle/pull/635).
|
||||
- Refactor the VarBuilder api
|
||||
[627](https://github.com/huggingface/candle/pull/627).
|
||||
- Add some quantization command
|
||||
[625](https://github.com/huggingface/candle/pull/625).
|
||||
- Support more quantized types, e.g. Q2K, Q4K, Q5K...
|
||||
[586](https://github.com/huggingface/candle/pull/586).
|
||||
- Add pose estimation to the yolo example
|
||||
[589](https://github.com/huggingface/candle/pull/589).
|
||||
- Api to write GGUF files
|
||||
[585](https://github.com/huggingface/candle/pull/585).
|
||||
- Support more quantization types
|
||||
[580](https://github.com/huggingface/candle/pull/580).
|
||||
- Add EfficientNet as an example Computer Vision model
|
||||
[572](https://github.com/huggingface/candle/pull/572).
|
||||
- Add a group parameter to convolutions
|
||||
[566](https://github.com/huggingface/candle/pull/566).
|
||||
- New dtype: int64
|
||||
[563](https://github.com/huggingface/candle/pull/563).
|
||||
- Handling of the GGUF file format.
|
||||
[559](https://github.com/huggingface/candle/pull/559).
|
||||
|
||||
## v0.1.2 - 2023-08-21
|
||||
Vendored
+82
@@ -0,0 +1,82 @@
|
||||
[workspace]
|
||||
members = [
|
||||
"candle-core",
|
||||
"candle-datasets",
|
||||
"candle-examples",
|
||||
"candle-book",
|
||||
"candle-nn",
|
||||
"candle-pyo3",
|
||||
"candle-transformers",
|
||||
"candle-wasm-examples/*",
|
||||
"candle-wasm-tests",
|
||||
"tensor-tools",
|
||||
]
|
||||
exclude = [
|
||||
"candle-flash-attn",
|
||||
"candle-kernels",
|
||||
"candle-metal-kernels",
|
||||
"candle-onnx",
|
||||
]
|
||||
resolver = "2"
|
||||
|
||||
[workspace.package]
|
||||
version = "0.8.3"
|
||||
edition = "2021"
|
||||
description = "Minimalist ML framework."
|
||||
repository = "https://github.com/huggingface/candle"
|
||||
keywords = ["blas", "tensor", "machine-learning"]
|
||||
categories = ["science"]
|
||||
license = "MIT OR Apache-2.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
ab_glyph = "0.2.23"
|
||||
accelerate-src = { version = "0.3.2" }
|
||||
anyhow = { version = "1", features = ["backtrace"] }
|
||||
byteorder = "1.4.3"
|
||||
candle = { path = "./candle-core", package = "candle-core", version = "0.8.3" }
|
||||
candle-datasets = { path = "./candle-datasets", version = "0.8.3" }
|
||||
candle-flash-attn = { path = "./candle-flash-attn", version = "0.8.3" }
|
||||
candle-kernels = { path = "./candle-kernels", version = "0.8.3" }
|
||||
candle-metal-kernels = { path = "./candle-metal-kernels", version = "0.8.3" }
|
||||
candle-nn = { path = "./candle-nn", version = "0.8.3" }
|
||||
candle-onnx = { path = "./candle-onnx", version = "0.8.3" }
|
||||
candle-transformers = { path = "./candle-transformers", version = "0.8.3" }
|
||||
clap = { version = "4.2.4", features = ["derive"] }
|
||||
criterion = { version = "0.5.1", default-features=false }
|
||||
cudarc = { version = "0.13.5", features = ["std", "cublas", "cublaslt", "curand", "driver", "nvrtc", "f16", "cuda-version-from-build-system", "dynamic-linking"], default-features=false }
|
||||
fancy-regex = "0.13.0"
|
||||
gemm = { version = "0.17.0", features = ["wasm-simd128-enable"] }
|
||||
hf-hub = "0.4.1"
|
||||
half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
|
||||
hound = "3.5.1"
|
||||
image = { version = "0.25.2", default-features = false, features = ["jpeg", "png"] }
|
||||
imageproc = { version = "0.24.0", default-features = false }
|
||||
intel-mkl-src = { version = "0.8.1", features = ["mkl-static-lp64-iomp"] }
|
||||
libc = { version = "0.2.147" }
|
||||
log = "0.4"
|
||||
memmap2 = { version = "0.9.3", features = ["stable_deref_trait"] }
|
||||
num_cpus = "1.15.0"
|
||||
num-traits = "0.2.15"
|
||||
parquet = { version = "51.0.0" }
|
||||
rand = "0.8.5"
|
||||
rand_distr = "0.4.3"
|
||||
rayon = "1.7.0"
|
||||
safetensors = "0.4.1"
|
||||
serde = { version = "1.0.171", features = ["derive"] }
|
||||
serde_plain = "1.0.2"
|
||||
serde_json = "1.0.99"
|
||||
thiserror = "1"
|
||||
tokenizers = { version = "0.19.1", default-features = false }
|
||||
tracing = "0.1.37"
|
||||
tracing-chrome = "0.7.1"
|
||||
tracing-subscriber = "0.3.7"
|
||||
ug = "0.1.0"
|
||||
ug-cuda = "0.1.0"
|
||||
ug-metal = "0.1.0"
|
||||
yoke = { version = "0.7.2", features = ["derive"] }
|
||||
zip = { version = "1.1.1", default-features = false }
|
||||
metal = { version = "0.27.0", features = ["mps"]}
|
||||
|
||||
[profile.release-with-debug]
|
||||
inherits = "release"
|
||||
debug = true
|
||||
Vendored
+201
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
Vendored
+23
@@ -0,0 +1,23 @@
|
||||
Permission is hereby granted, free of charge, to any
|
||||
person obtaining a copy of this software and associated
|
||||
documentation files (the "Software"), to deal in the
|
||||
Software without restriction, including without
|
||||
limitation the rights to use, copy, modify, merge,
|
||||
publish, distribute, sublicense, and/or sell copies of
|
||||
the Software, and to permit persons to whom the Software
|
||||
is furnished to do so, subject to the following
|
||||
conditions:
|
||||
|
||||
The above copyright notice and this permission notice
|
||||
shall be included in all copies or substantial portions
|
||||
of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
|
||||
ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
|
||||
TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
|
||||
PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
|
||||
SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
|
||||
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
|
||||
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
|
||||
IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
|
||||
DEALINGS IN THE SOFTWARE.
|
||||
Vendored
+16
@@ -0,0 +1,16 @@
|
||||
.PHONY: clean-ptx clean test
|
||||
|
||||
clean-ptx:
|
||||
find target -name "*.ptx" -type f -delete
|
||||
echo "" > candle-kernels/src/lib.rs
|
||||
touch candle-kernels/build.rs
|
||||
touch candle-examples/build.rs
|
||||
touch candle-flash-attn/build.rs
|
||||
|
||||
clean:
|
||||
cargo clean
|
||||
|
||||
test:
|
||||
cargo test
|
||||
|
||||
all: test
|
||||
Vendored
+425
@@ -0,0 +1,425 @@
|
||||
# candle
|
||||
[](https://discord.gg/hugging-face-879548962464493619)
|
||||
[](https://crates.io/crates/candle-core)
|
||||
[](https://docs.rs/candle-core)
|
||||
[](https://github.com/huggingface/candle/blob/main/LICENSE-MIT)
|
||||
[](https://github.com/huggingface/candle/blob/main/LICENSE-APACHE)
|
||||
|
||||
Candle is a minimalist ML framework for Rust with a focus on performance (including GPU support)
|
||||
and ease of use. Try our online demos:
|
||||
[whisper](https://huggingface.co/spaces/lmz/candle-whisper),
|
||||
[LLaMA2](https://huggingface.co/spaces/lmz/candle-llama2),
|
||||
[T5](https://huggingface.co/spaces/radames/Candle-T5-Generation-Wasm),
|
||||
[yolo](https://huggingface.co/spaces/lmz/candle-yolo),
|
||||
[Segment
|
||||
Anything](https://huggingface.co/spaces/radames/candle-segment-anything-wasm).
|
||||
|
||||
## Get started
|
||||
|
||||
Make sure that you have [`candle-core`](https://github.com/huggingface/candle/tree/main/candle-core) correctly installed as described in [**Installation**](https://huggingface.github.io/candle/guide/installation.html).
|
||||
|
||||
Let's see how to run a simple matrix multiplication.
|
||||
Write the following to your `myapp/src/main.rs` file:
|
||||
```rust
|
||||
use candle_core::{Device, Tensor};
|
||||
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let device = Device::Cpu;
|
||||
|
||||
let a = Tensor::randn(0f32, 1., (2, 3), &device)?;
|
||||
let b = Tensor::randn(0f32, 1., (3, 4), &device)?;
|
||||
|
||||
let c = a.matmul(&b)?;
|
||||
println!("{c}");
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
`cargo run` should display a tensor of shape `Tensor[[2, 4], f32]`.
|
||||
|
||||
|
||||
Having installed `candle` with Cuda support, simply define the `device` to be on GPU:
|
||||
|
||||
```diff
|
||||
- let device = Device::Cpu;
|
||||
+ let device = Device::new_cuda(0)?;
|
||||
```
|
||||
|
||||
For more advanced examples, please have a look at the following section.
|
||||
|
||||
## Check out our examples
|
||||
|
||||
These online demos run entirely in your browser:
|
||||
- [yolo](https://huggingface.co/spaces/lmz/candle-yolo): pose estimation and
|
||||
object recognition.
|
||||
- [whisper](https://huggingface.co/spaces/lmz/candle-whisper): speech recognition.
|
||||
- [LLaMA2](https://huggingface.co/spaces/lmz/candle-llama2): text generation.
|
||||
- [T5](https://huggingface.co/spaces/radames/Candle-T5-Generation-Wasm): text generation.
|
||||
- [Phi-1.5, and Phi-2](https://huggingface.co/spaces/radames/Candle-Phi-1.5-Wasm): text generation.
|
||||
- [Segment Anything Model](https://huggingface.co/spaces/radames/candle-segment-anything-wasm): Image segmentation.
|
||||
- [BLIP](https://huggingface.co/spaces/radames/Candle-BLIP-Image-Captioning): image captioning.
|
||||
|
||||
We also provide a some command line based examples using state of the art models:
|
||||
|
||||
- [LLaMA v1, v2, and v3](./candle-examples/examples/llama/): general LLM, includes
|
||||
the SOLAR-10.7B variant.
|
||||
- [Falcon](./candle-examples/examples/falcon/): general LLM.
|
||||
- [Codegeex4](./candle-examples/examples/codegeex4-9b/): Code completion,code interpreter,web search,fuction calling,repository-level
|
||||
- [GLM4](./candle-examples/examples/glm4/): Open Multilingual Multimodal Chat LMs by THUDM
|
||||
- [Gemma v1 and v2](./candle-examples/examples/gemma/): 2b and 7b+/9b general LLMs from Google Deepmind.
|
||||
- [RecurrentGemma](./candle-examples/examples/recurrent-gemma/): 2b and 7b
|
||||
Griffin based models from Google that mix attention with a RNN like state.
|
||||
- [Phi-1, Phi-1.5, Phi-2, and Phi-3](./candle-examples/examples/phi/): 1.3b,
|
||||
2.7b, and 3.8b general LLMs with performance on par with 7b models.
|
||||
- [StableLM-3B-4E1T](./candle-examples/examples/stable-lm/): a 3b general LLM
|
||||
pre-trained on 1T tokens of English and code datasets. Also supports
|
||||
StableLM-2, a 1.6b LLM trained on 2T tokens, as well as the code variants.
|
||||
- [Mamba](./candle-examples/examples/mamba/): an inference only
|
||||
implementation of the Mamba state space model.
|
||||
- [Mistral7b-v0.1](./candle-examples/examples/mistral/): a 7b general LLM with
|
||||
better performance than all publicly available 13b models as of 2023-09-28.
|
||||
- [Mixtral8x7b-v0.1](./candle-examples/examples/mixtral/): a sparse mixture of
|
||||
experts 8x7b general LLM with better performance than a Llama 2 70B model with
|
||||
much faster inference.
|
||||
- [StarCoder](./candle-examples/examples/bigcode/) and
|
||||
[StarCoder2](./candle-examples/examples/starcoder2/): LLM specialized to code generation.
|
||||
- [Qwen1.5](./candle-examples/examples/qwen/): Bilingual (English/Chinese) LLMs.
|
||||
- [RWKV v5 and v6](./candle-examples/examples/rwkv/): An RNN with transformer level LLM
|
||||
performance.
|
||||
- [Replit-code-v1.5](./candle-examples/examples/replit-code/): a 3.3b LLM specialized for code completion.
|
||||
- [Yi-6B / Yi-34B](./candle-examples/examples/yi/): two bilingual
|
||||
(English/Chinese) general LLMs with 6b and 34b parameters.
|
||||
- [Quantized LLaMA](./candle-examples/examples/quantized/): quantized version of
|
||||
the LLaMA model using the same quantization techniques as
|
||||
[llama.cpp](https://github.com/ggerganov/llama.cpp).
|
||||
|
||||
<img src="https://github.com/huggingface/candle/raw/main/candle-examples/examples/quantized/assets/aoc.gif" width="600">
|
||||
|
||||
- [Stable Diffusion](./candle-examples/examples/stable-diffusion/): text to
|
||||
image generative model, support for the 1.5, 2.1, SDXL 1.0 and Turbo versions.
|
||||
|
||||
<img src="https://github.com/huggingface/candle/raw/main/candle-examples/examples/stable-diffusion/assets/stable-diffusion-xl.jpg" width="200">
|
||||
|
||||
- [Wuerstchen](./candle-examples/examples/wuerstchen/): another text to
|
||||
image generative model.
|
||||
|
||||
<img src="https://github.com/huggingface/candle/raw/main/candle-examples/examples/wuerstchen/assets/cat.jpg" width="200">
|
||||
|
||||
- [yolo-v3](./candle-examples/examples/yolo-v3/) and
|
||||
[yolo-v8](./candle-examples/examples/yolo-v8/): object detection and pose
|
||||
estimation models.
|
||||
|
||||
<img src="https://github.com/huggingface/candle/raw/main/candle-examples/examples/yolo-v8/assets/bike.od.jpg" width="200"><img src="https://github.com/huggingface/candle/raw/main/candle-examples/examples/yolo-v8/assets/bike.pose.jpg" width="200">
|
||||
- [segment-anything](./candle-examples/examples/segment-anything/): image
|
||||
segmentation model with prompt.
|
||||
|
||||
<img src="https://github.com/huggingface/candle/raw/main/candle-examples/examples/segment-anything/assets/sam_merged.jpg" width="200">
|
||||
|
||||
- [SegFormer](./candle-examples/examples/segformer/): transformer based semantic segmentation model.
|
||||
- [Whisper](./candle-examples/examples/whisper/): speech recognition model.
|
||||
- [EnCodec](./candle-examples/examples/encodec/): high-quality audio compression
|
||||
model using residual vector quantization.
|
||||
- [MetaVoice](./candle-examples/examples/metavoice/): foundational model for
|
||||
text-to-speech.
|
||||
- [Parler-TTS](./candle-examples/examples/parler-tts/): large text-to-speech
|
||||
model.
|
||||
- [T5](./candle-examples/examples/t5), [Bert](./candle-examples/examples/bert/),
|
||||
[JinaBert](./candle-examples/examples/jina-bert/) : useful for sentence embeddings.
|
||||
- [DINOv2](./candle-examples/examples/dinov2/): computer vision model trained
|
||||
using self-supervision (can be used for imagenet classification, depth
|
||||
evaluation, segmentation).
|
||||
- [VGG](./candle-examples/examples/vgg/),
|
||||
[RepVGG](./candle-examples/examples/repvgg): computer vision models.
|
||||
- [BLIP](./candle-examples/examples/blip/): image to text model, can be used to
|
||||
generate captions for an image.
|
||||
- [CLIP](./candle-examples/examples/clip/): multi-model vision and language
|
||||
model.
|
||||
- [TrOCR](./candle-examples/examples/trocr/): a transformer OCR model, with
|
||||
dedicated submodels for hand-writing and printed recognition.
|
||||
- [Marian-MT](./candle-examples/examples/marian-mt/): neural machine translation
|
||||
model, generates the translated text from the input text.
|
||||
- [Moondream](./candle-examples/examples/moondream/): tiny computer-vision model
|
||||
that can answer real-world questions about images.
|
||||
|
||||
Run them using commands like:
|
||||
```
|
||||
cargo run --example quantized --release
|
||||
```
|
||||
|
||||
In order to use **CUDA** add `--features cuda` to the example command line. If
|
||||
you have cuDNN installed, use `--features cudnn` for even more speedups.
|
||||
|
||||
There are also some wasm examples for whisper and
|
||||
[llama2.c](https://github.com/karpathy/llama2.c). You can either build them with
|
||||
`trunk` or try them online:
|
||||
[whisper](https://huggingface.co/spaces/lmz/candle-whisper),
|
||||
[llama2](https://huggingface.co/spaces/lmz/candle-llama2),
|
||||
[T5](https://huggingface.co/spaces/radames/Candle-T5-Generation-Wasm),
|
||||
[Phi-1.5, and Phi-2](https://huggingface.co/spaces/radames/Candle-Phi-1.5-Wasm),
|
||||
[Segment Anything Model](https://huggingface.co/spaces/radames/candle-segment-anything-wasm).
|
||||
|
||||
For LLaMA2, run the following command to retrieve the weight files and start a
|
||||
test server:
|
||||
```bash
|
||||
cd candle-wasm-examples/llama2-c
|
||||
wget https://huggingface.co/spaces/lmz/candle-llama2/resolve/main/model.bin
|
||||
wget https://huggingface.co/spaces/lmz/candle-llama2/resolve/main/tokenizer.json
|
||||
trunk serve --release --port 8081
|
||||
```
|
||||
And then head over to
|
||||
[http://localhost:8081/](http://localhost:8081/).
|
||||
|
||||
<!--- ANCHOR: useful_libraries --->
|
||||
|
||||
## Useful External Resources
|
||||
- [`candle-tutorial`](https://github.com/ToluClassics/candle-tutorial): A
|
||||
very detailed tutorial showing how to convert a PyTorch model to Candle.
|
||||
- [`candle-lora`](https://github.com/EricLBuehler/candle-lora): Efficient and
|
||||
ergonomic LoRA implementation for Candle. `candle-lora` has
|
||||
out-of-the-box LoRA support for many models from Candle, which can be found
|
||||
[here](https://github.com/EricLBuehler/candle-lora/tree/master/candle-lora-transformers/examples).
|
||||
- [`optimisers`](https://github.com/KGrewal1/optimisers): A collection of optimisers
|
||||
including SGD with momentum, AdaGrad, AdaDelta, AdaMax, NAdam, RAdam, and RMSprop.
|
||||
- [`candle-vllm`](https://github.com/EricLBuehler/candle-vllm): Efficient platform for inference and
|
||||
serving local LLMs including an OpenAI compatible API server.
|
||||
- [`candle-ext`](https://github.com/mokeyish/candle-ext): An extension library to Candle that provides PyTorch functions not currently available in Candle.
|
||||
- [`candle-coursera-ml`](https://github.com/vishpat/candle-coursera-ml): Implementation of ML algorithms from Coursera's [Machine Learning Specialization](https://www.coursera.org/specializations/machine-learning-introduction) course.
|
||||
- [`kalosm`](https://github.com/floneum/floneum/tree/master/interfaces/kalosm): A multi-modal meta-framework in Rust for interfacing with local pre-trained models with support for controlled generation, custom samplers, in-memory vector databases, audio transcription, and more.
|
||||
- [`candle-sampling`](https://github.com/EricLBuehler/candle-sampling): Sampling techniques for Candle.
|
||||
- [`gpt-from-scratch-rs`](https://github.com/jeroenvlek/gpt-from-scratch-rs): A port of Andrej Karpathy's _Let's build GPT_ tutorial on YouTube showcasing the Candle API on a toy problem.
|
||||
- [`candle-einops`](https://github.com/tomsanbear/candle-einops): A pure rust implementation of the python [einops](https://github.com/arogozhnikov/einops) library.
|
||||
- [`atoma-infer`](https://github.com/atoma-network/atoma-infer): A Rust library for fast inference at scale, leveraging FlashAttention2 for efficient attention computation, PagedAttention for efficient KV-cache memory management, and multi-GPU support. It is OpenAI api compatible.
|
||||
- [`llms-from-scratch-rs`](https://github.com/nerdai/llms-from-scratch-rs): A comprehensive Rust translation of the code from Sebastian Raschka's Build an LLM from Scratch book.
|
||||
|
||||
If you have an addition to this list, please submit a pull request.
|
||||
|
||||
<!--- ANCHOR_END: useful_libraries --->
|
||||
|
||||
<!--- ANCHOR: features --->
|
||||
|
||||
## Features
|
||||
|
||||
- Simple syntax, looks and feels like PyTorch.
|
||||
- Model training.
|
||||
- Embed user-defined ops/kernels, such as [flash-attention v2](https://github.com/huggingface/candle/blob/89ba005962495f2bfbda286e185e9c3c7f5300a3/candle-flash-attn/src/lib.rs#L152).
|
||||
- Backends.
|
||||
- Optimized CPU backend with optional MKL support for x86 and Accelerate for macs.
|
||||
- CUDA backend for efficiently running on GPUs, multiple GPU distribution via NCCL.
|
||||
- WASM support, run your models in a browser.
|
||||
- Included models.
|
||||
- Language Models.
|
||||
- LLaMA v1, v2, and v3 with variants such as SOLAR-10.7B.
|
||||
- Falcon.
|
||||
- StarCoder, StarCoder2.
|
||||
- Phi 1, 1.5, 2, and 3.
|
||||
- Mamba, Minimal Mamba
|
||||
- Gemma v1 2b and 7b+, v2 2b and 9b.
|
||||
- Mistral 7b v0.1.
|
||||
- Mixtral 8x7b v0.1.
|
||||
- StableLM-3B-4E1T, StableLM-2-1.6B, Stable-Code-3B.
|
||||
- Replit-code-v1.5-3B.
|
||||
- Bert.
|
||||
- Yi-6B and Yi-34B.
|
||||
- Qwen1.5, Qwen1.5 MoE.
|
||||
- RWKV v5 and v6.
|
||||
- Quantized LLMs.
|
||||
- Llama 7b, 13b, 70b, as well as the chat and code variants.
|
||||
- Mistral 7b, and 7b instruct.
|
||||
- Mixtral 8x7b.
|
||||
- Zephyr 7b a and b (Mistral-7b based).
|
||||
- OpenChat 3.5 (Mistral-7b based).
|
||||
- Text to text.
|
||||
- T5 and its variants: FlanT5, UL2, MADLAD400 (translation), CoEdit (Grammar correction).
|
||||
- Marian MT (Machine Translation).
|
||||
- Text to image.
|
||||
- Stable Diffusion v1.5, v2.1, XL v1.0.
|
||||
- Wurstchen v2.
|
||||
- Image to text.
|
||||
- BLIP.
|
||||
- TrOCR.
|
||||
- Audio.
|
||||
- Whisper, multi-lingual speech-to-text.
|
||||
- EnCodec, audio compression model.
|
||||
- MetaVoice-1B, text-to-speech model.
|
||||
- Parler-TTS, text-to-speech model.
|
||||
- Computer Vision Models.
|
||||
- DINOv2, ConvMixer, EfficientNet, ResNet, ViT, VGG, RepVGG, ConvNeXT,
|
||||
ConvNeXTv2, MobileOne, EfficientVit (MSRA), MobileNetv4, Hiera, FastViT.
|
||||
- yolo-v3, yolo-v8.
|
||||
- Segment-Anything Model (SAM).
|
||||
- SegFormer.
|
||||
- File formats: load models from safetensors, npz, ggml, or PyTorch files.
|
||||
- Serverless (on CPU), small and fast deployments.
|
||||
- Quantization support using the llama.cpp quantized types.
|
||||
|
||||
<!--- ANCHOR_END: features --->
|
||||
|
||||
## How to use
|
||||
|
||||
<!--- ANCHOR: cheatsheet --->
|
||||
Cheatsheet:
|
||||
|
||||
| | Using PyTorch | Using Candle |
|
||||
|------------|------------------------------------------|------------------------------------------------------------------|
|
||||
| Creation | `torch.Tensor([[1, 2], [3, 4]])` | `Tensor::new(&[[1f32, 2.], [3., 4.]], &Device::Cpu)?` |
|
||||
| Creation | `torch.zeros((2, 2))` | `Tensor::zeros((2, 2), DType::F32, &Device::Cpu)?` |
|
||||
| Indexing | `tensor[:, :4]` | `tensor.i((.., ..4))?` |
|
||||
| Operations | `tensor.view((2, 2))` | `tensor.reshape((2, 2))?` |
|
||||
| Operations | `a.matmul(b)` | `a.matmul(&b)?` |
|
||||
| Arithmetic | `a + b` | `&a + &b` |
|
||||
| Device | `tensor.to(device="cuda")` | `tensor.to_device(&Device::new_cuda(0)?)?` |
|
||||
| Dtype | `tensor.to(dtype=torch.float16)` | `tensor.to_dtype(&DType::F16)?` |
|
||||
| Saving | `torch.save({"A": A}, "model.bin")` | `candle::safetensors::save(&HashMap::from([("A", A)]), "model.safetensors")?` |
|
||||
| Loading | `weights = torch.load("model.bin")` | `candle::safetensors::load("model.safetensors", &device)` |
|
||||
|
||||
<!--- ANCHOR_END: cheatsheet --->
|
||||
|
||||
|
||||
## Structure
|
||||
|
||||
- [candle-core](./candle-core): Core ops, devices, and `Tensor` struct definition
|
||||
- [candle-nn](./candle-nn/): Tools to build real models
|
||||
- [candle-examples](./candle-examples/): Examples of using the library in realistic settings
|
||||
- [candle-kernels](./candle-kernels/): CUDA custom kernels
|
||||
- [candle-datasets](./candle-datasets/): Datasets and data loaders.
|
||||
- [candle-transformers](./candle-transformers): transformers-related utilities.
|
||||
- [candle-flash-attn](./candle-flash-attn): Flash attention v2 layer.
|
||||
- [candle-onnx](./candle-onnx/): ONNX model evaluation.
|
||||
|
||||
## FAQ
|
||||
|
||||
### Why should I use Candle?
|
||||
|
||||
Candle's core goal is to *make serverless inference possible*. Full machine learning frameworks like PyTorch
|
||||
are very large, which makes creating instances on a cluster slow. Candle allows deployment of lightweight
|
||||
binaries.
|
||||
|
||||
Secondly, Candle lets you *remove Python* from production workloads. Python overhead can seriously hurt performance,
|
||||
and the [GIL](https://www.backblaze.com/blog/the-python-gil-past-present-and-future/) is a notorious source of headaches.
|
||||
|
||||
Finally, Rust is cool! A lot of the HF ecosystem already has Rust crates, like [safetensors](https://github.com/huggingface/safetensors) and [tokenizers](https://github.com/huggingface/tokenizers).
|
||||
|
||||
|
||||
### Other ML frameworks
|
||||
|
||||
- [dfdx](https://github.com/coreylowman/dfdx) is a formidable crate, with shapes being included
|
||||
in types. This prevents a lot of headaches by getting the compiler to complain about shape mismatches right off the bat.
|
||||
However, we found that some features still require nightly, and writing code can be a bit daunting for non rust experts.
|
||||
|
||||
We're leveraging and contributing to other core crates for the runtime so hopefully both crates can benefit from each
|
||||
other.
|
||||
|
||||
- [burn](https://github.com/burn-rs/burn) is a general crate that can leverage multiple backends so you can choose the best
|
||||
engine for your workload.
|
||||
|
||||
- [tch-rs](https://github.com/LaurentMazare/tch-rs.git) Bindings to the torch library in Rust. Extremely versatile, but they
|
||||
bring in the entire torch library into the runtime. The main contributor of `tch-rs` is also involved in the development
|
||||
of `candle`.
|
||||
|
||||
### Common Errors
|
||||
|
||||
#### Missing symbols when compiling with the mkl feature.
|
||||
|
||||
If you get some missing symbols when compiling binaries/tests using the mkl
|
||||
or accelerate features, e.g. for mkl you get:
|
||||
```
|
||||
= note: /usr/bin/ld: (....o): in function `blas::sgemm':
|
||||
.../blas-0.22.0/src/lib.rs:1944: undefined reference to `sgemm_' collect2: error: ld returned 1 exit status
|
||||
|
||||
= note: some `extern` functions couldn't be found; some native libraries may need to be installed or have their path specified
|
||||
= note: use the `-l` flag to specify native libraries to link
|
||||
= note: use the `cargo:rustc-link-lib` directive to specify the native libraries to link with Cargo
|
||||
```
|
||||
or for accelerate:
|
||||
```
|
||||
Undefined symbols for architecture arm64:
|
||||
"_dgemm_", referenced from:
|
||||
candle_core::accelerate::dgemm::h1b71a038552bcabe in libcandle_core...
|
||||
"_sgemm_", referenced from:
|
||||
candle_core::accelerate::sgemm::h2cf21c592cba3c47 in libcandle_core...
|
||||
ld: symbol(s) not found for architecture arm64
|
||||
```
|
||||
|
||||
This is likely due to a missing linker flag that was needed to enable the mkl library. You
|
||||
can try adding the following for mkl at the top of your binary:
|
||||
```rust
|
||||
extern crate intel_mkl_src;
|
||||
```
|
||||
or for accelerate:
|
||||
```rust
|
||||
extern crate accelerate_src;
|
||||
```
|
||||
|
||||
#### Cannot run the LLaMA examples: access to source requires login credentials
|
||||
|
||||
```
|
||||
Error: request error: https://huggingface.co/meta-llama/Llama-2-7b-hf/resolve/main/tokenizer.json: status code 401
|
||||
```
|
||||
|
||||
This is likely because you're not permissioned for the LLaMA-v2 model. To fix
|
||||
this, you have to register on the huggingface-hub, accept the [LLaMA-v2 model
|
||||
conditions](https://huggingface.co/meta-llama/Llama-2-7b-hf), and set up your
|
||||
authentication token. See issue
|
||||
[#350](https://github.com/huggingface/candle/issues/350) for more details.
|
||||
|
||||
#### Missing cute/cutlass headers when compiling flash-attn
|
||||
|
||||
```
|
||||
In file included from kernels/flash_fwd_launch_template.h:11:0,
|
||||
from kernels/flash_fwd_hdim224_fp16_sm80.cu:5:
|
||||
kernels/flash_fwd_kernel.h:8:10: fatal error: cute/algorithm/copy.hpp: No such file or directory
|
||||
#include <cute/algorithm/copy.hpp>
|
||||
^~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
compilation terminated.
|
||||
Error: nvcc error while compiling:
|
||||
```
|
||||
[cutlass](https://github.com/NVIDIA/cutlass) is provided as a git submodule so you may want to run the following command to check it in properly.
|
||||
```bash
|
||||
git submodule update --init
|
||||
```
|
||||
|
||||
#### Compiling with flash-attention fails
|
||||
|
||||
```
|
||||
/usr/include/c++/11/bits/std_function.h:530:146: error: parameter packs not expanded with ‘...’:
|
||||
```
|
||||
|
||||
This is a bug in gcc-11 triggered by the Cuda compiler. To fix this, install a different, supported gcc version - for example gcc-10, and specify the path to the compiler in the NVCC_CCBIN environment variable.
|
||||
```
|
||||
env NVCC_CCBIN=/usr/lib/gcc/x86_64-linux-gnu/10 cargo ...
|
||||
```
|
||||
|
||||
#### Linking error on windows when running rustdoc or mdbook tests
|
||||
|
||||
```
|
||||
Couldn't compile the test.
|
||||
---- .\candle-book\src\inference\hub.md - Using_the_hub::Using_in_a_real_model_ (line 50) stdout ----
|
||||
error: linking with `link.exe` failed: exit code: 1181
|
||||
//very long chain of linking
|
||||
= note: LINK : fatal error LNK1181: cannot open input file 'windows.0.48.5.lib'
|
||||
```
|
||||
|
||||
Make sure you link all native libraries that might be located outside a project target, e.g., to run mdbook tests, you should run:
|
||||
|
||||
```
|
||||
mdbook test candle-book -L .\target\debug\deps\ `
|
||||
-L native=$env:USERPROFILE\.cargo\registry\src\index.crates.io-6f17d22bba15001f\windows_x86_64_msvc-0.42.2\lib `
|
||||
-L native=$env:USERPROFILE\.cargo\registry\src\index.crates.io-6f17d22bba15001f\windows_x86_64_msvc-0.48.5\lib
|
||||
```
|
||||
|
||||
#### Extremely slow model load time with WSL
|
||||
|
||||
This may be caused by the models being loaded from `/mnt/c`, more details on
|
||||
[stackoverflow](https://stackoverflow.com/questions/68972448/why-is-wsl-extremely-slow-when-compared-with-native-windows-npm-yarn-processing).
|
||||
|
||||
#### Tracking down errors
|
||||
|
||||
You can set `RUST_BACKTRACE=1` to be provided with backtraces when a candle
|
||||
error is generated.
|
||||
|
||||
#### CudaRC error
|
||||
|
||||
If you encounter an error like this one `called `Result::unwrap()` on an `Err` value: LoadLibraryExW { source: Os { code: 126, kind: Uncategorized, message: "The specified module could not be found." } }` on windows. To fix copy and rename these 3 files (make sure they are in path). The paths depend on your cuda version.
|
||||
`c:\Windows\System32\nvcuda.dll` -> `cuda.dll`
|
||||
`c:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\bin\cublas64_12.dll` -> `cublas.dll`
|
||||
`c:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\bin\curand64_10.dll` -> `curand.dll`
|
||||
@@ -0,0 +1 @@
|
||||
book
|
||||
+48
@@ -0,0 +1,48 @@
|
||||
[package]
|
||||
name = "candle-book"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description.workspace = true
|
||||
repository.workspace = true
|
||||
keywords.workspace = true
|
||||
categories.workspace = true
|
||||
license.workspace = true
|
||||
readme = "README.md"
|
||||
|
||||
[dependencies]
|
||||
accelerate-src = { workspace = true, optional = true }
|
||||
candle = { workspace = true }
|
||||
candle-datasets = { workspace = true }
|
||||
candle-nn = { workspace = true }
|
||||
candle-transformers = { workspace = true }
|
||||
candle-flash-attn = { workspace = true, optional = true }
|
||||
safetensors = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
num-traits = { workspace = true }
|
||||
intel-mkl-src = { workspace = true, optional = true }
|
||||
cudarc = { workspace = true, optional = true }
|
||||
half = { workspace = true, optional = true }
|
||||
image = { workspace = true, optional = true }
|
||||
anyhow = { workspace = true }
|
||||
tokio = "1.43.0"
|
||||
|
||||
[dev-dependencies]
|
||||
byteorder = { workspace = true }
|
||||
hf-hub = { workspace = true, features=["tokio"]}
|
||||
clap = { workspace = true }
|
||||
memmap2 = { workspace = true }
|
||||
rand = { workspace = true }
|
||||
tokenizers = { workspace = true, features = ["onig"] }
|
||||
tracing = { workspace = true }
|
||||
tracing-chrome = { workspace = true }
|
||||
tracing-subscriber = { workspace = true }
|
||||
# Necessary to disambiguate with tokio in wasm examples which are 1.28.1
|
||||
parquet = { workspace = true }
|
||||
image = { workspace = true }
|
||||
|
||||
[build-dependencies]
|
||||
anyhow = { workspace = true }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
[book]
|
||||
authors = ["Nicolas Patry"]
|
||||
language = "en"
|
||||
multilingual = false
|
||||
src = "src"
|
||||
title = "Candle Documentation"
|
||||
@@ -0,0 +1,6 @@
|
||||
# Introduction
|
||||
|
||||
{{#include ../../README.md:features}}
|
||||
|
||||
|
||||
This book will introduce step by step how to use `candle`.
|
||||
+28
@@ -0,0 +1,28 @@
|
||||
# Summary
|
||||
|
||||
[Introduction](README.md)
|
||||
|
||||
# User Guide
|
||||
|
||||
- [Installation](guide/installation.md)
|
||||
- [Hello World - MNIST](guide/hello_world.md)
|
||||
- [PyTorch cheatsheet](guide/cheatsheet.md)
|
||||
|
||||
# Reference Guide
|
||||
|
||||
- [Running a model](inference/inference.md)
|
||||
- [Using the hub](inference/hub.md)
|
||||
- [Error management](error_manage.md)
|
||||
- [Training](training/training.md)
|
||||
- [Simplified](training/simplified.md)
|
||||
- [MNIST](training/mnist.md)
|
||||
- [Fine-tuning]()
|
||||
- [Serialization]()
|
||||
- [Advanced Cuda usage]()
|
||||
- [Writing a custom kernel]()
|
||||
- [Porting a custom kernel]()
|
||||
- [Using MKL]()
|
||||
- [Creating apps]()
|
||||
- [Creating a WASM app]()
|
||||
- [Creating a REST api webserver]()
|
||||
- [Creating a desktop Tauri app]()
|
||||
@@ -0,0 +1 @@
|
||||
# Using MKL
|
||||
@@ -0,0 +1 @@
|
||||
# Creating apps
|
||||
@@ -0,0 +1 @@
|
||||
# Creating a desktop Tauri app
|
||||
@@ -0,0 +1 @@
|
||||
# Creating a REST api webserver
|
||||
@@ -0,0 +1 @@
|
||||
# Creating a WASM app
|
||||
@@ -0,0 +1 @@
|
||||
# Chapter 1
|
||||
@@ -0,0 +1 @@
|
||||
# Advanced Cuda usage
|
||||
@@ -0,0 +1 @@
|
||||
# Porting a custom kernel
|
||||
@@ -0,0 +1 @@
|
||||
# Writing a custom kernel
|
||||
@@ -0,0 +1,51 @@
|
||||
# Error management
|
||||
|
||||
You might have seen in the code base a lot of `.unwrap()` or `?`.
|
||||
If you're unfamiliar with Rust check out the [Rust book](https://doc.rust-lang.org/book/ch09-02-recoverable-errors-with-result.html)
|
||||
for more information.
|
||||
|
||||
What's important to know though, is that if you want to know *where* a particular operation failed
|
||||
You can simply use `RUST_BACKTRACE=1` to get the location of where the model actually failed.
|
||||
|
||||
Let's see on failing code:
|
||||
|
||||
```rust,ignore
|
||||
let x = Tensor::zeros((1, 784), DType::F32, &device)?;
|
||||
let y = Tensor::zeros((1, 784), DType::F32, &device)?;
|
||||
let z = x.matmul(&y)?;
|
||||
```
|
||||
|
||||
Will print at runtime:
|
||||
|
||||
```bash
|
||||
Error: ShapeMismatchBinaryOp { lhs: [1, 784], rhs: [1, 784], op: "matmul" }
|
||||
```
|
||||
|
||||
|
||||
After adding `RUST_BACKTRACE=1`:
|
||||
|
||||
|
||||
```bash
|
||||
Error: WithBacktrace { inner: ShapeMismatchBinaryOp { lhs: [1, 784], rhs: [1, 784], op: "matmul" }, backtrace: Backtrace [{ fn: "candle::error::Error::bt", file: "/home/nicolas/.cargo/git/checkouts/candle-5bb8ef7e0626d693/f291065/candle-core/src/error.rs", line: 200 }, { fn: "candle::tensor::Tensor::matmul", file: "/home/nicolas/.cargo/git/checkouts/candle-5bb8ef7e0626d693/f291065/candle-core/src/tensor.rs", line: 816 }, { fn: "myapp::main", file: "./src/main.rs", line: 29 }, { fn: "core::ops::function::FnOnce::call_once", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/core/src/ops/function.rs", line: 250 }, { fn: "std::sys_common::backtrace::__rust_begin_short_backtrace", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/sys_common/backtrace.rs", line: 135 }, { fn: "std::rt::lang_start::{{closure}}", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/rt.rs", line: 166 }, { fn: "core::ops::function::impls::<impl core::ops::function::FnOnce<A> for &F>::call_once", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/core/src/ops/function.rs", line: 284 }, { fn: "std::panicking::try::do_call", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/panicking.rs", line: 500 }, { fn: "std::panicking::try", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/panicking.rs", line: 464 }, { fn: "std::panic::catch_unwind", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/panic.rs", line: 142 }, { fn: "std::rt::lang_start_internal::{{closure}}", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/rt.rs", line: 148 }, { fn: "std::panicking::try::do_call", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/panicking.rs", line: 500 }, { fn: "std::panicking::try", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/panicking.rs", line: 464 }, { fn: "std::panic::catch_unwind", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/panic.rs", line: 142 }, { fn: "std::rt::lang_start_internal", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/rt.rs", line: 148 }, { fn: "std::rt::lang_start", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/rt.rs", line: 165 }, { fn: "main" }, { fn: "__libc_start_main" }, { fn: "_start" }] }
|
||||
```
|
||||
|
||||
Not super pretty at the moment, but we can see error occurred on `{ fn: "myapp::main", file: "./src/main.rs", line: 29 }`
|
||||
|
||||
|
||||
Another thing to note, is that since Rust is compiled it is not necessarily as easy to recover proper stacktraces
|
||||
especially in release builds. We're using [`anyhow`](https://docs.rs/anyhow/latest/anyhow/) for that.
|
||||
The library is still young, please [report](https://github.com/LaurentMazare/candle/issues) any issues detecting where an error is coming from.
|
||||
|
||||
## Cuda error management
|
||||
|
||||
When running a model on Cuda, you might get a stacktrace not really representing the error.
|
||||
The reason is that CUDA is async by nature, and therefore the error might be caught while you were sending totally different kernels.
|
||||
|
||||
One way to avoid this is to use `CUDA_LAUNCH_BLOCKING=1` as an environment variable. This will force every kernel to be launched sequentially.
|
||||
You might still however see the error happening on other kernels as the faulty kernel might exit without an error but spoiling some pointer for which the error will happen when dropping the `CudaSlice` only.
|
||||
|
||||
|
||||
If this occurs, you can use [`compute-sanitizer`](https://docs.nvidia.com/compute-sanitizer/ComputeSanitizer/index.html)
|
||||
This tool is like `valgrind` but for cuda. It will help locate the errors in the kernels.
|
||||
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
# Pytorch cheatsheet
|
||||
|
||||
{{#include ../../../README.md:cheatsheet}}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user