Compare commits

...
Author SHA1 Message Date
Brandon Thomas 14feaccb91 Merge branch 'main' into feature/generatemode 2025-03-25 01:38:48 -04:00
Brandon Thomas 81c256e4f8 cleanup 2025-03-24 23:08:07 -04:00
Brandon Thomas 8d51d5b8ab can't trigger default features; switch to platform switches 2025-03-24 23:00:29 -04:00
Brandon Thomas 456bfd7bb9 trigger with --feature ml_models (need to make default) 2025-03-24 22:31:58 -04:00
Brandon Thomas cd3c864eff runs on mac 2025-03-24 21:46:25 -04:00
Brandon Thomas 990b7ff3e5 package split for mac (part one) (#459) 2025-03-24 20:18:57 -04:00
Brandon Thomas e1dbf4f2fe add inpaint demo/test endpoint 2025-03-22 03:16:48 -04:00
Tanish Manku d826db599d Remove the cancel generation button for now 2025-03-22 00:08:06 -07:00
Tanish Manku a642732c59 Clip overflow to have top rounded corners on image 2025-03-22 00:02:21 -07:00
Brandon Thomas 435c5aa0b4 fix strength bug 2025-03-22 02:57:06 -04:00
Tanish Manku 90b4aff846 Fix: Multiple invokes of a prompt after changing mode 2025-03-21 23:56:50 -07:00
Tanish Manku 4d3f097432 Hook up the TTI UI to invoke call 2025-03-21 23:38:25 -07:00
Tanish Manku e5de0e4731 Merge branch 'main' into feature/generatemode 2025-03-21 23:11:07 -07:00
Tanish Manku 04498b0457 Add initial Edit mode files 2025-03-21 23:02:17 -07:00
Brandon Thomas e3245dc593 fix on windows 2025-03-22 00:58:04 -04:00
bflatastic fbfc1e15f2 gallery gap fix 2025-03-21 22:46:30 +07:00
bflatastic 4a65317729 rename image to realtime 2025-03-21 20:13:13 +07:00
Tanish Manku e7ed0693a3 Add generation engine and signals 2025-03-21 00:44:44 -07:00
Tanish Manku 99ac8faf33 Merge branch 'main' into feature/generatemode 2025-03-20 23:41:19 -07:00
Brandon Thomas 08590d0c61 log environment details 2025-03-21 01:11:34 -04:00
Brandon Thomas 32cb39a643 text_to_image endpoint 2025-03-20 23:28:59 -04:00
bflatastic f0213d4a63 cancel button 2025-03-21 00:03:15 +07:00
Brandon Thomas 5fdb4d1024 disable msi builds due to 2.0 gb limit for 32bit 2025-03-20 11:59:43 -04:00
bflatastic 58f684b1e0 tab disabled prop and disabled message to block tab change when generating 2025-03-20 22:42:36 +07:00
bflatastic e83389a874 background gallery + many other styling and animation transition updates 2025-03-20 22:11:12 +07:00
Brandon Thomas d11945b685 download DIS model 2025-03-20 11:07:34 -04:00
Brandon Thomas f3babbcb19 if statements -> switch statements 2025-03-20 10:53:34 -04:00
Tanish Manku 3602f8fc2b Add generation signals 2025-03-20 02:10:27 -07:00
Tanish Manku 115c4fc60c Add a generation mode UI component 2025-03-19 23:58:53 -07:00
Brandon Thomas f37011ffc6 remove background: frontend UX (#458) 2025-03-20 01:41:43 -04:00
Brandon Thomas de683b750a remove background from images 2025-03-18 22:10:18 -04:00
Brandon Thomas 4603c65d2a fix strength slider range 2025-03-18 18:21:39 -04:00
Brandon Thomas c5b0d837a9 fix bug with float strengths from frontend 2025-03-18 01:48:54 -04:00
Brandon Thomas 652785224a easier windows builds. put dlls in the windows_dlls folder 2025-03-14 04:22:18 -04:00
Brandon Thomas 8d3ed947d1 change default prompt, log end of endpoint logic 2025-03-14 03:20:46 -04:00
Brandon Thomas 243583f0b7 kill more dead code 2025-03-13 23:39:04 -04:00
Brandon Thomas 150223c48c refactor weights (#455) 2025-03-13 23:29:20 -04:00
bflatastic ddf6a7f54e top center ui bug fix and padding 2025-03-13 22:28:26 +07:00
Brandon Thomas 17c6e549b5 ignore windows dlls 2025-03-13 00:38:10 -04:00
Brandon Thomas 031d9d8f20 Download models from Cloudflare instead of HF (#454) 2025-03-13 00:32:43 -04:00
Brandon Thomas 43df6c9181 app default size 2025-03-12 21:58:27 -04:00
Tanish Manku a8e029497f Fix context menu showing up after transforming (#453) 2025-03-12 15:18:00 -07:00
Brandon Thomas 4936662ddd correct default path on windows 2025-03-12 18:08:23 -04:00
BFlat 883e32220c tabs and app mode signals (#450)
Usage example - Conditional Rendering Based on App Mode:
```typescript
import { appMode } from "~/signals";

function ConditionalComponent() {
  return (
    <div className="app-container">
      {/* Show/hide components based on current mode */}
      {appMode.value === 'image' && (
        <div className="image-tools">
          <h3>Image Mode Tools</h3>
        </div>
      )}
      
      {appMode.value === 'edit' && (
        <div className="edit-tools">
          <h3>Edit Mode Tools</h3>
        </div>
      )}
  );
}
```
2025-03-12 13:26:10 -04:00
Brandon Thomas 727eead5a7 fix windows search for missing stdc++.lib during build 2025-03-12 01:27:27 -04:00
Kasisnu b4cc5b84cf f16 2025-03-11 22:30:02 +05:30
Brandon Thomas d7e77775ca Add new premium plan identifiers 2025-03-11 02:15:46 -04:00
Brandon Thomas b75472c434 fix windows builds by disabling flash attn/cutlass, openssl, etc. 2025-03-11 00:49:30 -04:00
Brandon Thomas a4f876daa9 Multithreaded downloader + notifications (#447) 2025-03-10 21:51:07 -04:00
Tanish Manku 4b309a5dd8 Render shapes as vectors and revamp rendering logic (#442) 2025-03-07 14:41:11 -08:00
Tanish Manku 08384955e4 Remove 'React-Eyedrop' and unused imports (#443)
As simple as it sounds
2025-03-06 23:05:23 -08:00
Brandon Thomas 31037b83a3 downloader thread 2025-03-06 22:48:48 -05:00
Brandon Thomas 53d93cb1ac change name of app 2025-03-06 21:45:01 -05:00
Brandon Thomas a0de8fc603 load all models from model registry 2025-03-06 21:39:19 -05:00
Brandon Thomas ef3527dd85 weights directory, assets directory 2025-03-06 21:06:58 -05:00
Brandon Thomas 1a887d7414 don't download anything from huggingface 2025-03-06 19:27:22 -05:00
Brandon Thomas 48abc0d731 no longer download clip encoders from huggingface 2025-03-06 15:22:10 -05:00
CrossPr0duct 0233105393 Merge branch 'main' of github.com:storytold/storyteller-rust 2025-03-06 15:12:28 -05:00
CrossPr0duct 8f94f81db4 Background Hack Implemented without engine refactors. 2025-03-06 15:12:04 -05:00
Kasisnu 3f25aefc3f starting point (#440) 2025-03-06 10:41:08 -05:00
BFlat c2929cfdb6 color picker ux and make paint cursor visible on both dark/light bg (#441)
ux updates to color picker and paint cursor
2025-03-06 20:52:25 +07:00
Brandon Thomas a859bd427b download clip models (not using them yet) 2025-03-06 08:42:44 -05:00
Brandon Thomas 706becfd7e download unet 2025-03-05 19:32:41 -05:00
Brandon Thomas ada1f46e62 model downloading (wip) 2025-03-05 18:24:05 -05:00
bflatastic 3556dcbfd3 removed annoying focus on buttons 2 2025-03-05 23:40:33 +07:00
bflatastic 4ab6ca1caa disabled annoying focus ring 2025-03-05 23:21:11 +07:00
BFlat 74f5b14d18 added loading indicator text (#439)
added signal for loading indicator text change
2025-03-05 10:23:48 +07:00
CrossPr0duct ac6dcfbc05 Selector Square Bug fixed.
Added Loading to RealTimeDrawEngine.
2025-03-04 20:20:33 -05:00
Brandon Thomas 68b0acf8ed always send event for now 2025-03-04 19:43:30 -05:00
Brandon Thomas f783910c0a vendor in cutlass 2025-03-04 19:28:28 -05:00
Brandon Thomas f3ee145f65 vendor in candle-examples 2025-03-04 19:18:18 -05:00
Brandon Thomas f8a2f4f858 marking function async makes tauri not lock up main thread 2025-03-04 18:28:28 -05:00
Brandon Thomas 3d7e558884 also send model type with event 2025-03-04 18:23:21 -05:00
CrossPr0duct 51d941a35d Merge rust 2025-03-04 18:01:24 -05:00
CrossPr0duct e56ce47205 Remove eraser button 2025-03-04 18:00:55 -05:00
Brandon Thomas 26c7b29c17 cleanup unet caching, report download 2025-03-04 17:51:14 -05:00
Brandon Thomas 89d2898e5f send notification events 2025-03-04 17:10:05 -05:00
Brandon Thomas 8124cbbbe1 destructure 2025-03-04 16:32:45 -05:00
Brandon Thomas f6b3ceb537 fix bad abstraction for vae caching 2025-03-04 16:24:28 -05:00
CrossPr0duct 606d1d2345 Added the progress bar. 2025-03-04 15:53:32 -05:00
Brandon Thomas e088f8e4fd send test event to javascript from tauri 2025-03-04 15:13:55 -05:00
Brandon Thomas c29d6377a4 log strength 2025-03-03 20:28:52 -05:00
Brandon Thomas 494e1e6b2b release build profile, stop writing debug file at app startup 2025-03-03 15:35:05 -05:00
BFlat 4d31f702b4 toaster for notifications (#437)
added toaster for notifications and alerts using sonner library

import:
```
import { toast } from 'sonner';
```

basic usage just call:
```
toast('This is an normal toast');

toast.error('This is an error toast');

toast.success('This is an success toast');

toast('This is an action toast', {
          action: {
            label: 'Action',
            onClick: () => console.log('Action!'),
          },
        });

toast('My cancel toast', {
  cancel: <Button onClick={() => console.log('Cancel!')}>Cancel</Button>,
});
```

more usage check here: https://sonner.emilkowal.ski/toast
2025-03-04 00:27:15 +07:00
Brandon Thomas 202b7876b4 temp test for windows 2025-03-02 20:31:34 -05:00
CrossPr0duct 360c3a268b Fixing the flow feels more natural,
swaps to shape selection or image selection when that is added,
except for render copy.
2025-03-01 21:25:27 -05:00
CrossPr0duct e2d9600458 Ensure the drag layer sits on top. 2025-03-01 21:06:21 -05:00
CrossPr0duct 31194c6622 Fixed brush issues and fixed issues with the paint server embedding lengths. 2025-03-01 19:10:16 -05:00
CrossPr0duct 5045e44af7 Painted img2img requires rescaling the image. 2025-03-01 18:09:19 -05:00
CrossPr0duct d387a04d12 Server Changes modified So it doesn't crash.
Type Fixes.
2025-03-01 17:46:29 -05:00
CrossPr0duct 0276349b5d Add Background color wire. 2025-03-01 17:03:14 -05:00
CrossPr0duct 3c037b085b Update the server lora strength.
Add Rendering for strength and text events.
2025-03-01 15:41:39 -05:00
ArEnScandCrossPr0duct 7965afd96c Feature/live drawing (#436)
Co-authored-by: CrossPr0duct <alicextensor@gmail.com>
2025-03-01 15:27:06 -05:00
BFlat 17dd093a23 model selection modal with signals, dummy data and dummy download progress (#435) 2025-03-01 11:39:42 -05:00
BFlat 989d36dc6c loading indicator signals (#434)
added signals for the loading indicator
2025-03-01 18:07:40 +07:00
BFlat d23199c046 added loading bar (dev only) (#433) 2025-02-27 23:48:58 +07:00
Kasisnu e67cacc23f basic lcm implementation 2025-02-27 22:00:58 +05:30
CrossPr0duct 5d9358103c Brush Size added in. 2025-02-26 20:29:22 -05:00
CrossPr0duct c36bfeb47c Merge branch 'main' of github.com:storytold/storyteller-rust 2025-02-26 19:25:46 -05:00
CrossPr0duct ad289122e0 Add Copy Canvas. 2025-02-26 19:25:15 -05:00
Tanish Manku 35a8920db8 Show context menu at right click position (#432) 2025-02-26 15:39:47 -08:00
Tanish Manku 701d84602f Add an eyedrop to all color pickers (#431) 2025-02-26 15:39:34 -08:00
CrossPr0duct 1d2a52556d Remove loading bar in context menu 2025-02-26 14:16:08 -05:00
bflatastic af44de74bc text prompt area to force blur on outside click 2025-02-26 23:17:15 +07:00
bflatastic 920cb9e751 slider center fix 2025-02-26 22:12:02 +07:00
BFlat c9c0a1bdf7 dialog transitions, model buttons, ratio button (#430)
**changes:**
- refined dialogs to use a new BaseDialog component which has sleek
transitions on open/close
- added model buttons on top center (only visible if in dev mode
currently) - these open up a modal and has no added signals yet.
- added ratio button on top right (uses new reusable ToggleButton
component) (only visible if in dev mode currently) - not signaled yet.
- added tooltip to background color button
2025-02-26 17:24:59 +07:00
Tanish Manku c65b00adc8 Add Background button and menu to toolbar (#427) 2025-02-26 00:47:51 -08:00
Tanish Manku e199086cff Context Menu Redesign (#425)
- Fix right click
- New design
2025-02-26 00:42:33 -08:00
Brandon Thomas 92d1195abf Vendor candle (#429) 2025-02-26 02:30:42 -05:00
BFlat d663eed493 render bug fix (#428) 2025-02-26 14:27:33 +07:00
CrossPr0duct 7c70634746 Merge branch 'main' of github.com:storytold/storyteller-rust 2025-02-25 23:12:57 -05:00
CrossPr0duct f6039c9967 Preview Copy still WIP UX 2025-02-25 23:12:33 -05:00
CrossPr0duct 8fea9c1a5e Working on previews 2025-02-25 23:09:48 -05:00
Tanish Manku ecee0d2ea8 Make color picker and paint menu use hex alpha picker (#426) 2025-02-25 18:12:41 -08:00
CrossPr0duct 0d41a52d5e Hide server stuff 2025-02-25 20:33:30 -05:00
CrossPr0duct 8ce8d111eb Merge branch 'main' of github.com:storytold/storyteller-rust 2025-02-25 18:40:12 -05:00
CrossPr0duct b1c20c7d42 Adding a test server. 2025-02-25 18:39:48 -05:00
Tanish Manku 000638ab9f Add brush size to Paint and Erase (#424)
- Both modes' menus have been separated for better maintenance
- Slider is generic, can be reused
- New signals added for engine to hook into
- Sliders tested working on the UI side
2025-02-25 02:45:52 -08:00
CrossPr0duct ef8b308cf0 Fixed issue with the translate and undo. 2025-02-24 23:37:20 -05:00
CrossPr0duct 5ea12fd682 Fix shape recurson due to app mode. 2025-02-24 22:41:16 -05:00
ArEnScandCrossPr0duct 1d5c4df906 Feature/paint (#423)
Paint Feature.
Fixes any issues involving rendering to the other frame for undo redo
stack.
Adds the painting feature as well as colors for it.
Fixes issues with transform updates.

---------

Co-authored-by: CrossPr0duct <alicextensor@gmail.com>
2025-02-24 21:11:33 -05:00
ae6434e11e Add mode switching: SELECT, PAINT, ERASE (#420)
Co-authored-by: BFlat <bombayapirat@gmail.com>
Co-authored-by: CrossPr0duct <alicextensor@gmail.com>
2025-02-23 21:30:00 -05:00
Michael Chung 6077232cd3 Hide Preview Drag Layer as it is not done. 2025-02-21 17:21:43 -05:00
Michael Chung f013373180 Copy Shift Still needs some UX work but it does work.
Need to store it as a image node.
2025-02-21 00:39:59 -05:00
Michael Chung f0c1fcf516 Got the color changes to work. 2025-02-20 22:34:17 -05:00
Brandon Thomas 9f392de9c5 bump packages to enable windows builds 2025-02-20 18:10:07 -05:00
BFlat f5fb340de9 landing hero update (#419) 2025-02-21 01:49:28 +07:00
Tanish Manku 0392bda03c Forward color picker events to engine (#418) 2025-02-20 11:38:29 -05:00
Michael Chung 1bca478621 Move color picker into the correct toolbar 2025-02-19 23:56:45 -05:00
Michael Chung f9c966003a Fixed issue with shapes not having the context menu component. 2025-02-19 23:18:28 -05:00
Michael Chung 2d8245cd55 Added Shapes.
Removed Extra buttons.
2025-02-19 22:29:31 -05:00
Tanish Manku 9d1d605f5c Add shapes to global toolbar and color picker to contextual (#417)
Shapes scaffolding is fully set up, color picker isn't tied up to the
nodes so it doesn't actually change the color yet
2025-02-19 21:55:06 -05:00
Michael Chung 1e38fbad5c Turns out you need to convert these rasters into images. 2025-02-19 20:33:01 -05:00
Michael Chung 6e03df968d Remove duplicate adds?
Then fixed layering issue by removing useless background layer.
Need to fix shape rendering.
2025-02-19 19:34:06 -05:00
Michael Chung 13e0e4bb02 Merge branch 'main' of github.com:storytold/storyteller-rust 2025-02-19 18:00:26 -05:00
Michael Chung 442c3579e1 Fixed Modal to have the image enter the scene right away.
Added Shape Fixes resizes properly and drags.
Removing Progress Fill.
2025-02-19 17:59:14 -05:00
bflatastic b8a305e271 mobile download button disable landing 2025-02-19 22:58:22 +07:00
BFlat a70061bfd3 realtime product landing (#415) 2025-02-19 20:44:22 +07:00
Brandon Thomas e7a392b8e2 add cfg (#416) 2025-02-19 00:21:24 -05:00
Michael Chung 9147da22f2 ShapeNodes work need to fix transformations. 2025-02-19 00:09:47 -05:00
Brandon Thomas 29c2df3ee6 strength, negative prompt 2025-02-18 21:42:22 -05:00
Brandon Thomas 1db4278712 cleanup 2025-02-18 02:19:38 -05:00
Brandon Thomas 5ee910ad07 pass prompt from frontend 2025-02-18 01:57:37 -05:00
Brandon ThomasandKasisnu ff9d2c6b34 realtime + speedup (#414)
Co-authored-by: Kasisnu <5044120+kasisnu@users.noreply.github.com>
2025-02-18 01:29:56 -05:00
Michael Chung 462dcd7645 Removed buttons that didn't need to be there.
Wired up prompt and strength sliders and save image.
Removed some more code.
2025-02-14 17:35:48 -05:00
Michael Chung d1617f88a4 Fixed UI Positioning.
Tied events to the system.
Should accept realtime inputs on changes.
2025-02-14 17:10:14 -05:00
ArEnScandTanish c2e57d0815 Prompt Box and Strength (#412)
Co-authored-by: Tanish <tanish2k09.dev@gmail.com>
2025-02-14 15:44:15 -05:00
6942 changed files with 1750746 additions and 1414 deletions
+4
View File
@@ -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/
+3
View File
@@ -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
+3137 -246
View File
File diff suppressed because it is too large Load Diff
+13
View File
@@ -23,6 +23,8 @@ members = [
# Desktop apps
"crates/desktop/src-tauri",
"crates/desktop/ml_models",
"crates/desktop/ml_weights_registry",
# Shared libraries
"crates/lib/actix_cors_configs",
@@ -96,6 +98,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",
+8
View File
@@ -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
+79
View File
@@ -0,0 +1,79 @@
[package]
name = "ml_models"
edition = "2021"
version = "0.0.1"
authors = ["storyteller"]
publish = false
description = ""
license = ""
repository = ""
[lib]
name = "ml_models"
path = "src/lib.rs"
crate-type = ["staticlib", "cdylib", "rlib"]
[package.metadata.cargo-udeps.ignore]
normal = ["my-workspace-hack"]
[dependencies]
# Internal
filesys = { path = "../../lib/files/filesys" }
memory_caching = { path = "../../lib/caching/memory_caching" }
ml_weights_registry = { path = "../ml_weights_registry" }
primitives = { path = "../../lib/data/primitives" }
# Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332
cudarc = "0.13.7"
# External
anyhow = "1.0"
base64 = "0.22.1"
bytes = "1.10"
candle-core = { path = "../../vendor/candle/candle-core", features = ["cuda"] }
candle-nn = { path = "../../vendor/candle/candle-nn", features = ["cuda"] }
candle-transformers = { path = "../../vendor/candle/candle-transformers", features = ["cuda"] }
const_format = "0.2.34"
directories = "6.0.0"
dirs = "6.0.0"
futures = "0.3"
half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
hf-hub = "0.3.2"
image = "0.25.5"
log = "0.4"
ndarray = "0.15.6"
nvidia-checker = "0.1.2"
once_cell = "1.20.3"
ort = { version = "1.16.2", features = ["cuda"] }
rand = "0.8.5"
# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls
reqwest = { version = "0.12.12", features = ["stream"], default-features = false }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
serde_yaml = "0.9"
tauri = { version = "2.2.4", features = [] }
tauri-plugin-log = "2"
tempdir = "0.3.7"
tempfile = "3.18.0"
tokenizers = "0.15.2"
tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] }
tokio-util = { version = "0.7", features = ["codec"] }
# tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] }
[features]
default = ["cuda"]
cuda = [
"candle-core/cuda",
"candle-nn/cuda",
"candle-transformers/cuda",
# NB(bt,2025-03-10): Building without Flash attention on Windows for now to bypass linker issues
"candle-transformers/flash-attn",
"candle-core/cudnn"
]
[dev-dependencies]
libloading = "0.8.6"
serial_test = "3.2.0"
assertor = "0.0.3"
+1
View File
@@ -0,0 +1 @@
pub mod ml;
@@ -0,0 +1,2 @@
pub mod onnx_session;
pub mod remove_image_background;
@@ -0,0 +1,19 @@
//! Adapted from https://github.com/dnanhkhoa/rust-background-removal/blob/main/src/onnx.rs
use std::path::Path;
use ort::{Environment, ExecutionProvider, GraphOptimizationLevel, SessionBuilder};
pub fn onnx_session(onnx_model_file: &Path) -> Result<ort::Session, anyhow::Error> {
let environment = Environment::default().into_arc();
let session = SessionBuilder::new(&environment)?
.with_optimization_level(GraphOptimizationLevel::Level1)? // Configure model optimization level
.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,61 @@
use crate::ml::background_removal::onnx_session::onnx_session;
use image::{imageops, DynamicImage};
use ndarray::{Array, CowArray};
use ort::Value;
use std::path::Path;
pub async fn remove_image_background<P: AsRef<Path>>(model_path: P, image: DynamicImage) -> anyhow::Result<DynamicImage> {
// NB: Real time needed cuDNN 9, but this uses cuDNN 8:
// sudo apt-get install libcudnn8-dev libcudnn8
// update-alternatives: warning: forcing reinstallation of alternative /usr/include/x86_64-linux-gnu/cudnn_v9.h because link group libcudnn is broken
// update-alternatives: using /usr/include/x86_64-linux-gnu/cudnn_v8.h to provide /usr/include/cudnn.h (libcudnn) in manual mode
let session = onnx_session(model_path.as_ref())?;
let image = process_dynamic_image(&session, image)?;
Ok(image)
}
fn process_dynamic_image(
session: &ort::Session,
dynamic_img: DynamicImage,
) -> Result<DynamicImage, anyhow::Error> {
let input_shape = session.inputs[0]
.dimensions()
.map(|dim| dim.unwrap())
.collect::<Vec<usize>>();
let input_img = dynamic_img.into_rgba8();
let scaling_factor = f32::min(
1., // Avoid upscaling
f32::min(
input_shape[3] as f32 / input_img.width() as f32, // Width ratio
input_shape[2] as f32 / input_img.height() as f32, // Height ratio
),
);
let mut resized_img = imageops::resize(
&input_img,
input_shape[3] as u32,
input_shape[2] as u32,
imageops::FilterType::Triangle,
);
let input_tensor = CowArray::from(
Array::from_shape_fn(input_shape, |indices| {
let mean = 128.;
let std = 256.;
(resized_img[(indices[3] as u32, indices[2] as u32)][indices[1]] as f32 - mean) / std
})
.into_dyn(),
);
let inputs = vec![Value::from_array(session.allocator(), &input_tensor)?];
let outputs = session.run(inputs)?;
let output_tensor = outputs[0].try_extract::<f32>()?;
for (indices, alpha) in output_tensor.view().indexed_iter() {
resized_img[(indices[3] as u32, indices[2] as u32)][3] = (alpha * 255.) as u8;
}
let output_img = imageops::resize(
&resized_img,
(input_img.width() as f32 * scaling_factor) as u32,
(input_img.height() as f32 * scaling_factor) as u32,
imageops::FilterType::Triangle,
);
Ok(DynamicImage::ImageRgba8(output_img))
}
@@ -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"))
}
+7
View File
@@ -0,0 +1,7 @@
pub mod background_removal;
pub mod image;
pub mod model_cache;
pub mod model_config;
pub mod model_file;
pub mod prompt_cache;
pub mod stable_diffusion;
@@ -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,61 @@
use crate::ml::model_file::StableDiffusionVersion;
use candle_core::{DType, Device};
use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
use hf_hub::api::sync::Api;
use log::info;
const DEFAULT_SD_IMAGE_WIDTH: usize = 1024;
const DEFAULT_SD_IMAGE_HEIGHT: usize = 1024;
const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32;
const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::V1_5;
pub struct ModelConfig {
/// We only support one device type currently, and we can't multiplex
/// over several devices yet.
pub device: Device,
/// The core numeric type to use for models.
pub dtype: DType,
pub sd_version: StableDiffusionVersion,
pub sd_config: StableDiffusionConfig,
/// Probably shouldn't be here.
pub hf_api: Api,
}
impl ModelConfig {
pub fn init() -> anyhow::Result<Self> {
let device = Device::new_cuda(0)
.unwrap_or_else(|e| {
info!("CUDA not available ({}), falling back to CPU", e);
Device::Cpu
});
let sd_config = match SD_VERSION {
StableDiffusionVersion::Turbo => {
StableDiffusionConfig::sdxl_turbo(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
}
StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => {
StableDiffusionConfig::v1_5(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
}
StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint => {
StableDiffusionConfig::v2_1(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
}
StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint => {
StableDiffusionConfig::sdxl(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
}
};
let hf_api = Api::new()?;
Ok(Self {
device,
dtype: DEFAULT_SD_NUMERIC_DATATYPE,
sd_version: SD_VERSION,
sd_config,
hf_api,
})
}
}
@@ -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,158 @@
#[cfg(feature = "accelerate")]
extern crate accelerate_src;
#[cfg(feature = "mkl")]
extern crate intel_mkl_src;
use candle_transformers::models::stable_diffusion;
use std::path::{Path, PathBuf};
use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
use ml_weights_registry::weights_registry::weights::{CLIP_JSON, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2};
use anyhow::Error as E;
use candle_core::{DType, Device, Module, Tensor, D};
use log::info;
use tokenizers::Tokenizer;
pub struct ClipArgs<'a, P: AsRef<Path>, Q: AsRef<Path>> {
pub prompt: &'a str,
pub uncond_prompt: &'a str,
pub tokenizer: Option<String>,
pub clip_weights: Option<String>,
pub clip2_weights: Option<String>,
pub sd_version: StableDiffusionVersion,
pub sd_config: &'a stable_diffusion::StableDiffusionConfig,
pub use_f16: bool,
pub device: &'a Device,
pub dtype: DType,
pub use_guide_scale: bool,
pub clip_json_path: P,
pub clip_weights_path: Q,
}
#[allow(clippy::too_many_arguments)]
pub fn infer_clip_text_embeddings<P: AsRef<Path>, Q: AsRef<Path>>(args: ClipArgs<'_, P, Q>) -> anyhow::Result<Tensor> {
let ClipArgs {
prompt,
uncond_prompt,
tokenizer,
clip_weights,
clip2_weights,
sd_version,
sd_config,
use_f16,
device,
dtype,
use_guide_scale,
clip_json_path,
clip_weights_path,
} = args;
info!("Preparing text embeddings... for {:?}", sd_version);
let which = match sd_version {
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,
clip_json_path.as_ref(),
clip_weights_path.as_ref(),
)
})
.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,
clip_json_path: &Path,
clip_weights_path: &Path,
) -> anyhow::Result<Tensor> {
info!("do_infer_clip_text_embeddings called with args {:?}", (prompt, uncond_prompt, tokenizer.clone().unwrap_or_else(|| "None".to_string()), clip_weights.clone().unwrap_or_else(|| "None".to_string()), clip2_weights.clone().unwrap_or_else(|| "None".to_string()), sd_version, sd_config, use_f16, device, dtype, use_guide_scale, first));
let tokenizer = clip_json_path;
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.");
// TODO(bt,2025-03-23): These used to be two different clip models. We should only need one.
let clip_weights = if first {
println!(">>>>> CLIP 1");
clip_weights_path
} else {
println!(">>>>> CLIP 2");
clip_weights_path
};
let clip_config = if first { &sd_config.clip } else { sd_config.clip2.as_ref().unwrap() };
let text_model = stable_diffusion::build_clip_transformer(clip_config, clip_weights, device, DType::F32)?;
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,281 @@
use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
use anyhow::{anyhow, Error as E, Result};
use candle_core::{DType, Device, IndexOp, Tensor, D};
use candle_nn::VarBuilder;
use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
use candle_transformers::models::stable_diffusion::lcm::{LCMScheduler, LCMSchedulerConfig};
use candle_transformers::models::stable_diffusion::unet_2d::BlockConfig;
use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModelConfig;
use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
use crate::ml::model_cache::ModelCache;
use crate::ml::model_config::ModelConfig;
use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
use crate::ml::prompt_cache::PromptCache;
use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
use crate::ml::stable_diffusion::infer_clip_text_embeddings::{infer_clip_text_embeddings, ClipArgs};
use crate::ml::stable_diffusion::remap_lcm_strength_range::remap_lcm_strength_range;
use image::DynamicImage;
use log::info;
use ml_weights_registry::weights_registry::weights::{LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
use rand::Rng;
const DEFAULT_STRENGTH : f64 = 75.0;
pub struct Args<'a, P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>> {
pub image: &'a DynamicImage,
pub prompt: String,
pub uncond_prompt: String,
pub cfg_scale: Option<f64>,
pub img2img_strength: Option<f64>,
pub model_cache: &'a ModelCache,
pub prompt_cache: &'a PromptCache,
pub use_flash_attn: bool,
pub model_config: &'a ModelConfig,
pub maybe_seed: Option<u64>,
pub scheduler_steps: usize,
pub vae_path: P1,
pub unet_path: P2,
pub clip_json_path: P3,
pub clip_weights_path: P4,
}
pub fn lcm_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<RgbImage>
where P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>
{
let Args {
prompt,
uncond_prompt,
cfg_scale,
img2img_strength,
model_cache,
prompt_cache,
image,
use_flash_attn ,
model_config: ModelConfig {
device,
dtype,
sd_version,
sd_config,
hf_api,
},
maybe_seed,
scheduler_steps,
vae_path,
unet_path,
clip_json_path,
clip_weights_path,
} = args;
let dtype = *dtype;
let sd_version = *sd_version;
let img2img_strength = img2img_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(scheduler_steps, img2img_strength, LCMSchedulerConfig::default())?;
let seed = maybe_seed.unwrap_or_else(|| rand::thread_rng().gen());
let seed = 42;
device.set_seed(seed)?;
let guidance_scale = match cfg_scale {
Some(guidance_scale) => guidance_scale,
None => match 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!("Prompt is NOT cached! Calculating embedding...");
let tensor = infer_clip_text_embeddings(ClipArgs {
prompt: & prompt,
uncond_prompt: &uncond_prompt,
tokenizer: None, // tokenizer
clip_weights: None,
clip2_weights: None, // clip2_weights
sd_version,
sd_config: &sd_config,
use_f16, // use_f16
device,
dtype,
use_guide_scale,
clip_json_path,
clip_weights_path,
})?;
prompt_cache.store_copy(&prompt, &tensor)?;
tensor
};
info!("Text embeddings shape: {:?}", text_embeddings.shape());
info!("Loading input image into tensor...");
let input_image = dynamic_image_to_tensor(image, &device, dtype)?;
info!("Reference image shape: {:?}", input_image.shape());
let vae_scale = get_vae_scale(sd_version);
let maybe_vae = model_cache.get_vae()?;
let vae = match maybe_vae {
Some(vae) => vae,
None => {
info!("Building VAE model from file {:?}...", vae_path.as_ref());
let vae = sd_config.build_vae(vae_path, &device, dtype)?;
let vae = Arc::new(vae);
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 in_channels = match 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_path], dtype, &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, &device, dtype)?;
// Generate latents from input image using the LCM approach
info!("Generating latents from input image...");
let init_latents = (init_latent_dist.sample()? * vae_scale)?;
info!("Initial latents shape: {:?}", init_latents.shape());
// Calculate timesteps for LCM with proper img2img handling
let timesteps = LCMScheduler::get_timesteps_for_steps(scheduler_steps, img2img_strength);
let t_start = timesteps[0]; // First denoising step
// Add noise at the appropriate timestep for img2img
info!("Adding noise to latents for t_start={}", t_start);
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, &timestep) in timesteps.iter().enumerate() {
info!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep);
let latent_model_input = if use_guide_scale { Tensor::cat(&[&latents, &latents], 0)? } else { latents.clone() };
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) => {
info!("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, scheduler_steps)?;
}
info!("Diffusion process completed, decoding image...");
let image = vae.decode(&(latents / vae_scale)?)?;
info!("VAE decode completed");
info!("Scaling image back");
let image = ((image / 2.)? + 0.5)?;
info!("Normalized image values");
let image = (image.clamp(0f32, 1.)? * 255.)?;
info!("Convert to int8");
let image = image.to_dtype(DType::U8)?;
info!("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,270 @@
use std::path::Path;
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
use crate::ml::model_cache::ModelCache;
use crate::ml::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, ClipArgs};
use anyhow::{anyhow, Error as E, Result};
use candle_core::{DType, Device, IndexOp, Tensor, D};
use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
use image::DynamicImage;
use log::info;
use rand::Rng;
pub struct Args<'a, P: AsRef<Path>, Q: AsRef<Path>, R: AsRef<Path>, S: AsRef<Path>> {
pub image: &'a DynamicImage,
pub prompt: String,
pub uncond_prompt: String,
pub cfg_scale: Option<f64>,
pub i2i_strength: Option<u8>,
pub model_cache: &'a ModelCache,
pub prompt_cache: &'a PromptCache,
pub sd_version: StableDiffusionVersion,
pub sd_config: StableDiffusionConfig,
pub device: &'a Device,
pub dtype: DType,
pub maybe_seed: Option<u64>,
pub scheduler_steps: usize,
pub sdxl_turbo_vae_path: P,
pub sdxl_turbo_unet_path: Q,
pub clip_json_path: R,
pub clip_weights_path: S,
}
pub fn stable_diffusion_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<RgbImage>
where P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>
{
let Args {
prompt,
uncond_prompt,
cfg_scale,
i2i_strength,
model_cache,
prompt_cache,
image,
sd_version,
sd_config,
device,
dtype,
maybe_seed,
scheduler_steps,
sdxl_turbo_vae_path,
sdxl_turbo_unet_path,
clip_json_path,
clip_weights_path,
} = args;
println!("Starting image generation with the following configuration:");
println!(" Model: {:?}", sd_version);
println!(" Prompt: {}", prompt);
println!(" Steps: {}", scheduler_steps);
println!(" Device: {:?}", device);
println!("Model dimensions: {}x{}", sd_config.width, sd_config.height);
// TODO(bt,2025-02-18): The scheduler is `EulerAncestralDiscreteScheduler`, but we may want to port an LCM scheduler.
// This is a target for performance improvement
// See: https://github.com/huggingface/candle/issues/1331
let mut scheduler = sd_config.build_scheduler(scheduler_steps)?;
let seed = maybe_seed.unwrap_or_else(|| rand::thread_rng().gen());
device.set_seed(seed)?;
info!("Using seed: {}", seed);
let guidance_scale = match cfg_scale {
Some(guidance_scale) => guidance_scale,
None => match 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(ClipArgs {
prompt: &prompt,
uncond_prompt: &uncond_prompt,
tokenizer: None, // tokenizer
clip_weights: None, // clip_weights
clip2_weights: None, // clip2_weights
sd_version: sd_version,
sd_config: &sd_config,
use_f16: false, // use_f16
device: device,
dtype: dtype,
use_guide_scale,
clip_json_path: clip_json_path.as_ref(),
clip_weights_path: clip_weights_path.as_ref(),
})?;
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, device, dtype)?;
println!("Reference image shape: {:?}", input_image.shape());
let vae_scale = get_vae_scale(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 = sdxl_turbo_vae_path.as_ref();
let vae = sd_config.build_vae(vae_file, &device, 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 = sdxl_turbo_unet_path.as_ref();
let unet = sd_config.build_unet(unet_weights_file, &device, 4, true, dtype)
.map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
let unet = Arc::new(unet);
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 = scheduler_steps - (scheduler_steps as f64 * img2img_strength) as usize;
println!("Starting from step {} of {} (strength: {})", start, 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, &timestep) 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,18 @@
[package]
name = "ml_weights_registry"
edition = "2021"
version = "0.0.1"
authors = ["storyteller"]
publish = false
description = ""
license = ""
repository = ""
[lib]
name = "ml_weights_registry"
path = "src/lib.rs"
# External
[dependencies]
anyhow = "1.0"
const_format = "0.2.34"
@@ -0,0 +1 @@
pub mod weights_registry;
@@ -0,0 +1,3 @@
mod weight_descriptor_builder;
pub mod weight_descriptor;
pub mod weights;
@@ -0,0 +1,45 @@
// 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 {
ImageSegmentation,
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::ImageSegmentation => "image segmentation",
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::weights_registry::weight_descriptor::WeightDescriptor {
name: $name,
filename: $file_name,
r2_download_url: const_format::concatcp!(crate::weights_registry::weight_descriptor_builder::R2_BUCKET_URL, $file_name),
function: $weight_function,
}
};
}
pub(crate) use weight;
@@ -0,0 +1,65 @@
use crate::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction};
use crate::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
);
/// Image segmentation and background removal
/// dis_medium.onnx
pub const DIS_MEDIUM_ONNX : WeightDescriptor = weight!(
"Dichotomous Image Segmentation Medium",
"dis_medium.onnx",
WeightFunction::ImageSegmentation
);
#[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
);
+4
View File
@@ -0,0 +1,4 @@
# Don't restart the dev server on Rust changes
crates/
Cargo.toml
Cargo.lock
+60 -7
View File
@@ -7,6 +7,7 @@ license = ""
repository = ""
edition = "2021"
rust-version = "1.77.2"
publish = false
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -18,11 +19,63 @@ 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" }
ml_weights_registry = { path = "../ml_weights_registry" }
primitives = { path = "../../lib/data/primitives" }
# Internal ML
# TODO: Only on Windows and Linux
#ml_models = { path = "../ml_models", optional = true }
# External
#openssl = { version = "0.10", features = ["vendored"] }
anyhow = "1.0"
base64 = "0.22.1"
bytes = "1.10"
const_format = "0.2.34"
directories = "6.0.0"
dirs = "6.0.0"
futures = "0.3"
half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
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"
tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] }
tokio-util = { version = "0.7", features = ["codec"] }
# tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] }
[target.'cfg(not(target_os = "macos"))'.dependencies]
ml_models = { path = "../ml_models" }
[dev-dependencies]
libloading = "0.8.6"
serial_test = "3.2.0"
assertor = "0.0.3"
#[target.'cfg(target_os = "macos")'.features]
#default = []
#ml_models = ["dep:ml_models"]
[profile.dev]
incremental = true # Compile your binary in smaller steps.
[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.
Binary file not shown.

After

Width:  |  Height:  |  Size: 8.4 KiB

@@ -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()))
}
@@ -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,24 @@
use crate::ml::model_cache::ModelCache;
use crate::ml::prompt_cache::PromptCache;
use crate::state::app_config::AppConfig;
use crate::state::app_dir::AppDataRoot;
use image::{ImageReader, RgbImage};
use log::info;
use tauri::{AppHandle, State};
#[tauri::command]
pub async fn inpaint_image(
image: &str,
mask: &str,
prompt: String,
model_config: State<'_, AppConfig>,
model_cache: State<'_, ModelCache>,
prompt_cache: State<'_, PromptCache>,
app_data_root: State<'_, AppDataRoot>,
app: AppHandle,
) -> Result<String, String> {
info!("inpaint_endpoint endpoint called.");
Ok(mask.to_string())
}
@@ -0,0 +1,6 @@
pub mod download_models;
pub mod flip_image;
pub mod inpaint_image_endpoint;
pub mod realtime_image_endpoint;
pub mod remove_background_endpoint;
pub mod text_to_image_endpoint;
@@ -0,0 +1,143 @@
#[cfg(not(target_os = "macos"))]
use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
use crate::state::app_config::AppConfig;
use crate::state::app_dir::AppDataRoot;
use crate::stubs::model_cache::ModelCache;
use crate::stubs::prompt_cache::PromptCache;
use crate::utils::image::decode_base64_image::decode_base64_image;
use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png;
use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_png;
use image::imageops::FilterType;
use image::{DynamicImage, RgbImage};
use log::{error, info};
use ml_weights_registry::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
use std::path::PathBuf;
use tauri::{AppHandle, State};
const PROMPT_FILENAME : &str = "prompt.txt";
const PROMPT: &str = "shiba inu, cute dog, detailed, walking in a wooded forest, photorealistic, 8k";
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 mut image = image.resize(512, 512, FilterType::CatmullRom);
#[cfg(not(target_os = "macos"))]
{
let result = infer_image_impl(
&prompt,
image,
strength,
&model_config,
&model_cache,
prompt_cache,
&app_data_root
).await;
image = match result {
Ok(img) => DynamicImage::ImageRgb8(img),
Err(err) => {
error!("There was an error generating the image: {:?}", err);
return Err(format!("There was an error generating the image: {}", err));
}
};
}
let bytes = encode_dynamic_image_base64_png(image)
.map_err(|err| format!("failure to encode image: {:?}", err))?;
info!("Inference successful; image converted to base64, serving back to browser...");
Ok(bytes)
}
#[cfg(not(target_os = "macos"))]
async fn infer_image_impl(
prompt: &str,
image: DynamicImage,
strength: Option<f64>,
config: &AppConfig,
model_cache: &ModelCache,
prompt_cache: State<'_, PromptCache>,
app_data_root: &AppDataRoot,
) -> Result<RgbImage, String> {
let weights_dir = app_data_root.weights_dir();
let vae_path = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE);
let unet_path = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
let clip_json_path = weights_dir.weight_path(&CLIP_JSON);
//let clip_weights_path= weights_dir.weight_path(&LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16);
let clip_weights_path= weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER);
let args = Args {
image: &image,
prompt: prompt.to_string(),
uncond_prompt: NEGATIVE_PROMPT.to_string(),
model_cache,
prompt_cache: &prompt_cache,
img2img_strength: strength,
cfg_scale: config.cfg_scale,
use_flash_attn: true,
model_config: &config.model_config,
maybe_seed: None,
scheduler_steps: config.scheduler_steps,
vae_path: &vae_path,
unet_path: &unet_path,
clip_json_path: &clip_json_path,
clip_weights_path: &clip_weights_path,
};
let image = lcm_pipeline(args)
.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,45 @@
//! Adapted from https://github.com/dnanhkhoa/rust-background-removal
#[cfg(not(target_os = "macos"))]
use ml_models::ml::background_removal::remove_image_background::remove_image_background;
use crate::state::app_dir::AppDataRoot;
use crate::utils::image::decode_base64_image::decode_base64_image;
use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png;
use image::{imageops, DynamicImage};
use log::info;
use ml_weights_registry::weights_registry::weights::DIS_MEDIUM_ONNX;
use tauri::State;
/// This handler removes the background from an image.
#[tauri::command]
pub async fn remove_background(
image: &str,
app_data_root: State<'_, AppDataRoot>,
) -> Result<String, String> {
info!("remove_background endpoint called.");
let mut image = decode_base64_image(image)
.map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
#[cfg(not(target_os = "macos"))]
{
image = remove_background_impl(&app_data_root, image).await?;
}
let image = encode_dynamic_image_base64_png(image)
.map_err(|err| format!("failure to encode image: {:?}", err))?;
Ok(image)
}
#[cfg(not(target_os = "macos"))]
async fn remove_background_impl(app_data_root: &AppDataRoot, image: DynamicImage) -> Result<DynamicImage, String> {
let model_path = app_data_root.weights_dir().weight_path(&DIS_MEDIUM_ONNX);
let image = remove_image_background(model_path, image)
.await
.map_err(|err| format!("Couldn't remove image from image: {}", err))?;
Ok(image)
}
@@ -0,0 +1,115 @@
#[cfg(not(target_os = "macos"))]
use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
use crate::state::app_config::AppConfig;
use crate::state::app_dir::AppDataRoot;
use crate::stubs::model_cache::ModelCache;
use crate::stubs::prompt_cache::PromptCache;
use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png;
use image::imageops::FilterType;
use image::{DynamicImage, ImageReader, RgbImage};
use log::{error, info};
use ml_weights_registry::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
use once_cell::sync::Lazy;
use std::io::Cursor;
use tauri::{AppHandle, State};
const RANDOM_SEED: u32 = 42;
const STRENGTH: f64 = 100.0;
const PNG_BYTES : &[u8] = include_bytes!("../../binary_includes/1024.png");
// NB: We're just hacking a temporary "text to image" test endpoint leveraging the existing image to image modality.
static IMAGE : Lazy<DynamicImage> = Lazy::new(||{
ImageReader::new(Cursor::new(PNG_BYTES))
.with_guessed_format()
.expect("failed to detect png format")
.decode()
.expect("failed to decode image")
});
#[tauri::command]
pub async fn text_to_image(
prompt: String,
model_config: State<'_, AppConfig>,
model_cache: State<'_, ModelCache>,
prompt_cache: State<'_, PromptCache>,
app_data_root: State<'_, AppDataRoot>,
app: AppHandle,
) -> Result<String, String> {
info!("text_to_image endpoint called.");
let mut image = IMAGE.resize(512, 512, FilterType::CatmullRom);
#[cfg(not(target_os = "macos"))]
{
let result = text_to_image_impl(
&prompt,
image,
Some(STRENGTH),
&model_config,
&model_cache,
prompt_cache,
app,
&app_data_root
).await;
image = match result {
Ok(img) => DynamicImage::ImageRgb8(img),
Err(err) => {
error!("There was an error generating the image: {:?}", err);
return Err(format!("There was an error generating the image: {}", err));
}
};
}
let bytes = encode_dynamic_image_base64_png(image)
.map_err(|err| format!("failure to encode image: {:?}", err))?;
info!("Inference successful; image converted to base64, serving back to browser...");
Ok(bytes)
}
#[cfg(not(target_os = "macos"))]
async fn text_to_image_impl(
prompt: &str,
image: DynamicImage,
strength: Option<f64>,
config: &AppConfig,
model_cache: &ModelCache,
prompt_cache: State<'_, PromptCache>,
app: AppHandle,
app_data_root: &AppDataRoot,
) -> Result<RgbImage, String> {
let weights_dir = app_data_root.weights_dir();
let vae_path = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE);
let unet_path = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
let clip_json_path = weights_dir.weight_path(&CLIP_JSON);
let clip_weights_path= weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER);
let args = Args {
image: &image,
prompt: prompt.to_string(),
uncond_prompt: "".to_string(),
model_cache,
prompt_cache: &prompt_cache,
img2img_strength: strength,
cfg_scale: config.cfg_scale,
use_flash_attn: true,
model_config: &config.model_config,
maybe_seed: None,
scheduler_steps: config.scheduler_steps,
vae_path: &vae_path,
unet_path: &unet_path,
clip_json_path: &clip_json_path,
clip_weights_path: &clip_weights_path,
};
let image = lcm_pipeline(args)
.map_err(|err| format!("failure to encode image: {:?}", err))?;
Ok(image)
}
@@ -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 {
ImageSegmentation,
TextEncoder,
TextTokenizer,
Unet,
Vae,
}
#[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,
}
}
+67 -41
View File
@@ -1,55 +1,81 @@
mod image_endpoint;
pub mod endpoints;
pub mod events;
pub mod state;
pub mod stubs;
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::inpaint_image_endpoint::inpaint_image;
use crate::endpoints::flip_image::flip_image;
use crate::endpoints::realtime_image_endpoint::infer_image;
use crate::endpoints::remove_background_endpoint::remove_background;
use crate::endpoints::text_to_image_endpoint::text_to_image;
use crate::state::app_config::AppConfig;
use crate::stubs::model_cache::ModelCache;
use crate::stubs::prompt_cache::PromptCache;
use crate::threads::downloader_thread::downloader_thread;
use 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(),
)?;
}
// TODO(bt): This is broken on windows
// log_environment_details();
//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,
flip_image,
infer_image,
inpaint_image,
remove_background,
text_to_image,
])
.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
}
+9
View File
@@ -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,62 @@
use crate::state::app_dir::AppDataRoot;
use crate::state::os_platform::OsPlatform;
use crate::state::yaml_config::YamlConfig;
use crate::stubs::model_config::ModelConfig;
const DEFAULT_SD_IMAGE_WIDTH: usize = 1024;
const DEFAULT_SD_IMAGE_HEIGHT: usize = 1024;
/// A central place to configure app-wide details.
pub struct AppConfig {
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>,
/// Root location for application data.
pub app_data_root: AppDataRoot,
pub model_config: ModelConfig,
}
impl AppConfig {
pub fn init() -> anyhow::Result<Self> {
let yaml_configs = YamlConfig::load_from_config_file_recursive();
println!("Configs: {:?}", yaml_configs);
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 {
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,
app_data_root,
model_config: ModelConfig::init()?,
})
}
}
@@ -0,0 +1,187 @@
use crate::state::expanduser::expanduser;
use crate::state::os_platform::OsPlatform;
use anyhow::anyhow;
use directories::UserDirs;
use ml_weights_registry::weights_registry::weight_descriptor::WeightDescriptor;
use std::path::{Path, PathBuf};
use tempdir::TempDir;
use tempfile::{Builder, NamedTempFile};
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,3 @@
pub mod model_cache;
pub mod model_config;
pub mod prompt_cache;
@@ -0,0 +1,14 @@
#[cfg(not(target_os = "macos"))]
pub use ml_models::ml::model_cache::ModelCache;
#[cfg(not(not(target_os = "macos")))]
pub struct ModelCache {
// Intentionally left blank
}
#[cfg(not(not(target_os = "macos")))]
impl ModelCache {
pub fn new() -> Self {
Self {}
}
}
@@ -0,0 +1,14 @@
#[cfg(not(target_os = "macos"))]
pub use ml_models::ml::model_config::ModelConfig;
#[cfg(not(not(target_os = "macos")))]
pub struct ModelConfig {
// Intentionally left blank
}
#[cfg(not(not(target_os = "macos")))]
impl ModelConfig {
pub fn init() -> anyhow::Result<Self> {
Ok(Self {})
}
}
@@ -0,0 +1,15 @@
#[cfg(not(target_os = "macos"))]
pub use ml_models::ml::prompt_cache::PromptCache;
#[cfg(not(not(target_os = "macos")))]
pub struct PromptCache {
// Intentionally left blank
}
#[cfg(not(not(target_os = "macos")))]
impl PromptCache {
pub fn with_capacity(capacity: usize) -> anyhow::Result<Self> {
Ok(Self {})
}
}
@@ -0,0 +1,181 @@
use crate::events::notification_event::{NotificationEvent, NotificationModelType};
use crate::state::app_dir::AppDataRoot;
use crate::transfer::download::{download_async, ProgressUpdate};
use futures::StreamExt;
use futures::stream::FuturesUnordered;
use log::{debug, error, info};
use ml_weights_registry::weights_registry::weight_descriptor::WeightFunction;
use ml_weights_registry::weights_registry::weights::{CLIP_JSON, DIS_MEDIUM_ONNX, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2, SDXL_TURBO_UNET, SDXL_TURBO_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
use std::collections::VecDeque;
use std::path::PathBuf;
use std::sync::mpsc::Receiver;
use std::sync::{Arc, Mutex};
use tauri::{AppHandle, Emitter};
use tempfile::NamedTempFile;
use tokio::task::JoinHandle;
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(DIS_MEDIUM_ONNX);
//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,
WeightFunction::ImageSegmentation => NotificationModelType::ImageSegmentation,
};
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,16 @@
use base64::prelude::BASE64_STANDARD;
use base64::Engine;
use image::{ImageFormat, RgbImage};
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,11 @@
use log::info;
//use nvidia_checker::{get_cuda_version, get_cudnn_version, get_kernel_version, get_nvidia_driver_version, get_os_version, get_tensorrt_version};
pub fn log_environment_details() {
//info!("NVIDIA driver version: {}", get_nvidia_driver_version(false));
//info!("CUDA version: {}", get_cuda_version(false));
//info!("cuDNN version: {}", get_cudnn_version(false));
//info!("TensorRT version: {}", get_tensorrt_version(false));
//info!("OS version: {}", get_os_version(false));
//info!("Kernel version: {}", get_kernel_version(false));
}
@@ -0,0 +1,2 @@
pub mod image;
pub mod log_environment_details;
+4 -4
View File
@@ -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
}
@@ -0,0 +1,7 @@
#[test]
fn test_dynamic_library_loading() {
let name = libloading::library_filename("cublas");
assert_eq!(name, "libcublas.so");
}
+1
View File
@@ -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);
}
}
+1
View File
@@ -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
View File
@@ -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
View File
@@ -0,0 +1,7 @@
version: 2
updates:
- package-ecosystem: "cargo"
directory: "/"
schedule:
interval: "weekly"
open-pull-requests-limit: 5
+40
View File
@@ -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
View File
@@ -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/
+34
View File
@@ -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.
+68
View File
@@ -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
+87
View File
@@ -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
+15
View File
@@ -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
+48
View File
@@ -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
+3
View File
@@ -0,0 +1,3 @@
[submodule "candle-examples/examples/flash-attn/cutlass"]
path = candle-flash-attn/cutlass
url = https://github.com/NVIDIA/cutlass.git
+15
View File
@@ -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",
]
+113
View File
@@ -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
+82
View File
@@ -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
+201
View File
@@ -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.
+23
View File
@@ -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.
+16
View File
@@ -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
+425
View File
@@ -0,0 +1,425 @@
# candle
[![discord server](https://dcbadge.vercel.app/api/server/hugging-face-879548962464493619)](https://discord.gg/hugging-face-879548962464493619)
[![Latest version](https://img.shields.io/crates/v/candle-core.svg)](https://crates.io/crates/candle-core)
[![Documentation](https://docs.rs/candle-core/badge.svg)](https://docs.rs/candle-core)
[![License](https://img.shields.io/github/license/base-org/node?color=blue)](https://github.com/huggingface/candle/blob/main/LICENSE-MIT)
[![License](https://img.shields.io/badge/license-Apache%202.0-blue?style=flat-square)](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
View File
@@ -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
View File
@@ -0,0 +1,6 @@
[book]
authors = ["Nicolas Patry"]
language = "en"
multilingual = false
src = "src"
title = "Candle Documentation"

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