Added reference images to the generation

enqueueImageGeneration now accepts additionalImages enqueueImageGeneration needs testing. Polling for jobs doesn't work, most api calls don't work and need integration testing.
Merge branch 'main' into feature/gen
2026-10-28 20:42:50 -04:00 · 2025-04-03 17:26:52 -04:00 · 2025-04-03 12:37:31 -04:00 · 2025-04-03 23:24:14 +07:00 · 2025-04-03 23:01:02 +07:00 · 2025-04-03 11:55:38 -04:00
3195 changed files with 154862 additions and 2400 deletions
@@ -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/
@@ -559,15 +559,15 @@ checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0"

 [[package]]
 name = "android_log-sys"
-version = "0.3.1"
+version = "0.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5ecc8056bf6ab9892dcd53216c83d1597487d7dacac16c8df6b877d127df9937"
+checksum = "84521a3cf562bc62942e294181d9eef17eb38ceb8c68677bc49f144e4c3d4f8d"

 [[package]]
 name = "android_logger"
-version = "0.14.1"
+version = "0.15.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "05b07e8e73d720a1f2e4b6014766e6039fd2e96a4fa44e2a78d0e1fa2ff49826"
+checksum = "f6f39be698127218cca460cb624878c9aa4e2b47dba3b277963d2bf00bad263b"
 dependencies = [
 "android_log-sys",
 "env_filter",
@@ -658,23 +658,27 @@ name = "app"
 version = "0.1.0"
 dependencies = [
 "anyhow",
+ "assertor",
 "base64 0.22.1",
 "bytes 1.10.0",
- "candle-core",
- "candle-nn",
- "candle-transformers",
+ "chrono",
+ "const_format",
+ "crockford",
 "cudarc",
 "directories 6.0.0",
 "dirs 6.0.0",
 "filesys",
 "futures 0.3.29",
 "half",
- "hf-hub 0.3.2",
 "image",
 "libloading 0.8.6",
 "log",
 "memory_caching",
+ "ml_models",
+ "ml_weights_registry",
 "once_cell",
+ "ort",
+ "primitives",
 "rand 0.8.5",
 "reqwest 0.12.12",
 "serde",
@@ -686,7 +690,6 @@ dependencies = [
 "tauri-plugin-log",
 "tempdir",
 "tempfile",
- "tokenizers 0.15.2",
 "tokio",
 "tokio-util",
 ]
@@ -855,6 +858,15 @@ version = "1.1.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "10d2119f741b79fe9907f5396d19bffcb46568cfcc315e78677d731972ac7085"

+[[package]]
+name = "assertor"
+version = "0.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4b02f2af4042fbbed4d5732c517666629512c574c8e2edb5cfe2cf8a86564d52"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "async-channel"
 version = "1.9.0"
@@ -1365,7 +1377,7 @@ version = "0.5.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2c132eebf10f5cad5289222520a4a058514204aed6d791f1cf4fe8088b82d15f"
 dependencies = [
- "objc2",
+ "objc2 0.5.2",
 ]

 [[package]]
@@ -1827,6 +1839,7 @@ dependencies = [
 "fancy-regex",
 "half",
 "intel-mkl-src",
+ "log",
 "num-traits",
 "rand 0.8.5",
 "rayon",
@@ -1869,6 +1882,18 @@ dependencies = [
 "toml 0.8.14",
 ]

+[[package]]
+name = "carrot_utils"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3f34277d609ff503187aa8e8604913b103e2eb8f893222a4948748f22527252e"
+dependencies = [
+ "chrono",
+ "colored",
+ "log",
+ "simplelog",
+]
+
 [[package]]
 name = "cast"
 version = "0.3.0"
@@ -1951,9 +1976,9 @@ checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"

 [[package]]
 name = "chrono"
-version = "0.4.39"
+version = "0.4.40"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7e36cc9d416881d2e24f9a963be5fb1cd90966419ac844274161d10488b3e825"
+checksum = "1a7964611d71df112cb1730f2ee67324fcf4d0fc6606acbbe9bfe06df124637c"
 dependencies = [
 "android-tzdata",
 "iana-time-zone",
@@ -1961,7 +1986,7 @@ dependencies = [
 "num-traits",
 "serde",
 "wasm-bindgen",
- "windows-targets 0.52.6",
+ "windows-link",
 ]

 [[package]]
@@ -2046,9 +2071,9 @@ dependencies = [

 [[package]]
 name = "clap"
-version = "4.4.8"
+version = "4.5.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2275f18819641850fa26c89acc84d465c1bf91ce57bc2748b28c420473352f64"
+checksum = "0fbb260a053428790f3de475e304ff84cdbc4face759ea7a3e64c1edd938a7fc"
 dependencies = [
 "clap_builder",
 "clap_derive",
@@ -2056,23 +2081,23 @@ dependencies = [

 [[package]]
 name = "clap_builder"
-version = "4.4.8"
+version = "4.5.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "07cdf1b148b25c1e1f7a42225e30a0d99a615cd4637eae7365548dd4529b95bc"
+checksum = "64b17d7ea74e9f833c7dbf2cbe4fb12ff26783eda4782a8975b72f895c9b4d99"
 dependencies = [
 "anstream",
 "anstyle",
 "clap_lex",
- "strsim 0.10.0",
+ "strsim 0.11.1",
 ]

 [[package]]
 name = "clap_derive"
-version = "4.4.7"
+version = "4.5.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cf9804afaaf59a91e75b022a30fb7229a7901f60c755489cc61c9b423b836442"
+checksum = "501d359d5f3dcaf6ecdeee48833ae73ec6e42723a1e52419c79abf9507eec0a0"
 dependencies = [
- "heck 0.4.1",
+ "heck 0.5.0",
 "proc-macro2",
 "quote",
 "syn 2.0.87",
@@ -2080,9 +2105,9 @@ dependencies = [

 [[package]]
 name = "clap_lex"
-version = "0.6.0"
+version = "0.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "702fc72eb24e5a1e48ce58027a675bc24edd52096d5397d4aea7c6dd9eca0bd1"
+checksum = "f46ad14479a25103f283c0f10005961cf086d8dc42205bb44c46ac563475dca6"

 [[package]]
 name = "cloud_storage"
@@ -2175,6 +2200,16 @@ version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0b6a852b24ab71dffc585bcb46eaf7959d175cb865a7152e35b348d1b2960422"

+[[package]]
+name = "colored"
+version = "2.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c"
+dependencies = [
+ "lazy_static",
+ "windows-sys 0.59.0",
+]
+
 [[package]]
 name = "combine"
 version = "3.8.1"
@@ -2300,6 +2335,26 @@ dependencies = [
 "tiny-keccak",
 ]

+[[package]]
+name = "const_format"
+version = "0.2.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "126f97965c8ad46d6d9163268ff28432e8f6a1196a55578867832e3049df63dd"
+dependencies = [
+ "const_format_proc_macros",
+]
+
+[[package]]
+name = "const_format_proc_macros"
+version = "0.2.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d57c2eccfb16dbac1f4e61e206105db5820c9d26c3c472bc17c774259ef7744"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-xid",
+]
+
 [[package]]
 name = "constant_time_eq"
 version = "0.1.5"
@@ -3899,9 +3954,9 @@ checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"

 [[package]]
 name = "fdeflate"
-version = "0.3.1"
+version = "0.3.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "64d6dafc854908ff5da46ff3f8f473c6984119a2876a383a860246dd7841a868"
+checksum = "1e6853b52649d4ac5c0bd02320cddc5ba956bdb407c4b75a2c6b75bf51500f8c"
 dependencies = [
 "simd-adler32",
 ]
@@ -5408,9 +5463,9 @@ dependencies = [

 [[package]]
 name = "image"
-version = "0.25.5"
+version = "0.25.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cd6f44aed642f18953a158afeb30206f4d50da59fbc66ecb53c66488de73563b"
+checksum = "db35664ce6b9810857a38a906215e75a9c879f0696556a39f59c62829710251a"
 dependencies = [
 "bytemuck",
 "byteorder-lite",
@@ -5439,6 +5494,16 @@ dependencies = [
 "quick-error",
 ]

+[[package]]
+name = "image_algorithms"
+version = "0.0.1"
+dependencies = [
+ "errors",
+ "image",
+ "my-workspace-hack",
+ "testing",
+]
+
 [[package]]
 name = "imageproc"
 version = "0.24.0"
@@ -6257,9 +6322,9 @@ checksum = "9374ef4228402d4b7e403e5838cb880d9ee663314b0a900d5a6aabf0c213552e"

 [[package]]
 name = "log"
-version = "0.4.21"
+version = "0.4.27"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "90ed8c1e510134f979dbc4f070f87d4313098b704861a105fe34231c70a3901c"
+checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94"
 dependencies = [
 "value-bag",
 ]
@@ -6679,6 +6744,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b3b1c9bd4fe1f0f8b387f6eb9eb3b4a1aa26185e5750efb9140301703f62cd1b"
 dependencies = [
 "adler2",
+ "simd-adler32",
 ]

 [[package]]
@@ -6747,6 +6813,59 @@ dependencies = [
 "ws2_32-sys",
 ]

+[[package]]
+name = "ml_models"
+version = "0.0.1"
+dependencies = [
+ "anyhow",
+ "assertor",
+ "base64 0.22.1",
+ "bytes 1.10.0",
+ "candle-core",
+ "candle-nn",
+ "candle-transformers",
+ "const_format",
+ "cudarc",
+ "directories 6.0.0",
+ "dirs 6.0.0",
+ "filesys",
+ "futures 0.3.29",
+ "half",
+ "hf-hub 0.3.2",
+ "image",
+ "image_algorithms",
+ "libloading 0.8.6",
+ "log",
+ "memory_caching",
+ "ml_weights_registry",
+ "ndarray",
+ "nvidia-checker",
+ "once_cell",
+ "ort",
+ "primitives",
+ "rand 0.8.5",
+ "reqwest 0.12.12",
+ "serde",
+ "serde_json",
+ "serde_yaml",
+ "serial_test",
+ "tauri",
+ "tauri-plugin-log",
+ "tempdir",
+ "tempfile",
+ "tokenizers 0.15.2",
+ "tokio",
+ "tokio-util",
+]
+
+[[package]]
+name = "ml_weights_registry"
+version = "0.0.1"
+dependencies = [
+ "anyhow",
+ "const_format",
+]
+
 [[package]]
 name = "mockall"
 version = "0.12.1"
@@ -6819,9 +6938,9 @@ dependencies = [
 "dpi",
 "gtk",
 "keyboard-types",
- "objc2",
+ "objc2 0.5.2",
 "objc2-app-kit",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 "once_cell",
 "png",
 "serde",
@@ -7011,6 +7130,19 @@ dependencies = [
 "tempfile",
 ]

+[[package]]
+name = "ndarray"
+version = "0.15.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb12d4e967ec485a5f71c6311fe28158e9d6f4bc4a447b474184d0f91a8fa32"
+dependencies = [
+ "matrixmultiply",
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "rawpointer",
+]
+
 [[package]]
 name = "ndk"
 version = "0.8.0"
@@ -7288,6 +7420,23 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"

+[[package]]
+name = "nvidia-checker"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c8b584e4b0be14d4e3f8b779d79477aa0f36658d7a535c42abc725013d0c8951"
+dependencies = [
+ "carrot_utils",
+ "chrono",
+ "clap",
+ "colored",
+ "dirs 5.0.1",
+ "log",
+ "serde",
+ "serde_derive",
+ "toml 0.8.14",
+]
+
 [[package]]
 name = "objc"
 version = "0.2.7"
@@ -7317,6 +7466,15 @@ dependencies = [
 "objc2-encode",
 ]

+[[package]]
+name = "objc2"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3531f65190d9cff863b77a99857e74c314dd16bf56c538c4b57c7cbc3f3a6e59"
+dependencies = [
+ "objc2-encode",
+]
+
 [[package]]
 name = "objc2-app-kit"
 version = "0.2.2"
@@ -7326,10 +7484,10 @@ dependencies = [
 "bitflags 2.5.0",
 "block2",
 "libc",
- "objc2",
+ "objc2 0.5.2",
 "objc2-core-data",
 "objc2-core-image",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 "objc2-quartz-core",
 ]

@@ -7341,9 +7499,9 @@ checksum = "74dd3b56391c7a0596a295029734d3c1c5e7e510a4cb30245f8221ccea96b009"
 dependencies = [
 "bitflags 2.5.0",
 "block2",
- "objc2",
+ "objc2 0.5.2",
 "objc2-core-location",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 ]

 [[package]]
@@ -7353,8 +7511,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a5ff520e9c33812fd374d8deecef01d4a840e7b41862d849513de77e44aa4889"
 dependencies = [
 "block2",
- "objc2",
- "objc2-foundation",
+ "objc2 0.5.2",
+ "objc2-foundation 0.2.2",
 ]

 [[package]]
@@ -7365,8 +7523,8 @@ checksum = "617fbf49e071c178c0b24c080767db52958f716d9eabdf0890523aeae54773ef"
 dependencies = [
 "bitflags 2.5.0",
 "block2",
- "objc2",
- "objc2-foundation",
+ "objc2 0.5.2",
+ "objc2-foundation 0.2.2",
 ]

 [[package]]
@@ -7376,8 +7534,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "55260963a527c99f1819c4f8e3b47fe04f9650694ef348ffd2227e8196d34c80"
 dependencies = [
 "block2",
- "objc2",
- "objc2-foundation",
+ "objc2 0.5.2",
+ "objc2-foundation 0.2.2",
 "objc2-metal",
 ]

@@ -7388,9 +7546,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "000cfee34e683244f284252ee206a27953279d370e309649dc3ee317b37e5781"
 dependencies = [
 "block2",
- "objc2",
+ "objc2 0.5.2",
 "objc2-contacts",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 ]

 [[package]]
@@ -7408,7 +7566,17 @@ dependencies = [
 "bitflags 2.5.0",
 "block2",
 "libc",
- "objc2",
+ "objc2 0.5.2",
+]
+
+[[package]]
+name = "objc2-foundation"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3a21c6c9014b82c39515db5b396f91645182611c97d24637cf56ac01e5f8d998"
+dependencies = [
+ "bitflags 2.5.0",
+ "objc2 0.6.0",
 ]

 [[package]]
@@ -7418,9 +7586,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a1a1ae721c5e35be65f01a03b6d2ac13a54cb4fa70d8a5da293d7b0020261398"
 dependencies = [
 "block2",
- "objc2",
+ "objc2 0.5.2",
 "objc2-app-kit",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 ]

 [[package]]
@@ -7431,8 +7599,8 @@ checksum = "dd0cba1276f6023976a406a14ffa85e1fdd19df6b0f737b063b95f6c8c7aadd6"
 dependencies = [
 "bitflags 2.5.0",
 "block2",
- "objc2",
- "objc2-foundation",
+ "objc2 0.5.2",
+ "objc2-foundation 0.2.2",
 ]

 [[package]]
@@ -7443,8 +7611,8 @@ checksum = "e42bee7bff906b14b167da2bac5efe6b6a07e6f7c0a21a7308d40c960242dc7a"
 dependencies = [
 "bitflags 2.5.0",
 "block2",
- "objc2",
- "objc2-foundation",
+ "objc2 0.5.2",
+ "objc2-foundation 0.2.2",
 "objc2-metal",
 ]

@@ -7454,8 +7622,8 @@ version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0a684efe3dec1b305badae1a28f6555f6ddd3bb2c2267896782858d5a78404dc"
 dependencies = [
- "objc2",
- "objc2-foundation",
+ "objc2 0.5.2",
+ "objc2-foundation 0.2.2",
 ]

 [[package]]
@@ -7466,12 +7634,12 @@ checksum = "b8bb46798b20cd6b91cbd113524c490f1686f4c4e8f49502431415f3512e2b6f"
 dependencies = [
 "bitflags 2.5.0",
 "block2",
- "objc2",
+ "objc2 0.5.2",
 "objc2-cloud-kit",
 "objc2-core-data",
 "objc2-core-image",
 "objc2-core-location",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 "objc2-link-presentation",
 "objc2-quartz-core",
 "objc2-symbols",
@@ -7486,8 +7654,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "44fa5f9748dbfe1ca6c0b79ad20725a11eca7c2218bceb4b005cb1be26273bfe"
 dependencies = [
 "block2",
- "objc2",
- "objc2-foundation",
+ "objc2 0.5.2",
+ "objc2-foundation 0.2.2",
 ]

 [[package]]
@@ -7498,9 +7666,9 @@ checksum = "76cfcbf642358e8689af64cee815d139339f3ed8ad05103ed5eaf73db8d84cb3"
 dependencies = [
 "bitflags 2.5.0",
 "block2",
- "objc2",
+ "objc2 0.5.2",
 "objc2-core-location",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 ]

 [[package]]
@@ -7511,9 +7679,9 @@ checksum = "68bc69301064cebefc6c4c90ce9cba69225239e4b8ff99d445a2b5563797da65"
 dependencies = [
 "bitflags 2.5.0",
 "block2",
- "objc2",
+ "objc2 0.5.2",
 "objc2-app-kit",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 ]

 [[package]]
@@ -7793,6 +7961,27 @@ dependencies = [
 "hashbrown 0.14.5",
 ]

+[[package]]
+name = "ort"
+version = "1.16.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "889dca4c98efa21b1ba54ddb2bde44fd4920d910f492b618351f839d8428d79d"
+dependencies = [
+ "flate2",
+ "half",
+ "lazy_static",
+ "libc",
+ "libloading 0.7.4",
+ "ndarray",
+ "tar",
+ "thiserror 1.0.63",
+ "tracing",
+ "ureq",
+ "vswhom",
+ "winapi 0.3.9",
+ "zip 0.6.6",
+]
+
 [[package]]
 name = "overload"
 version = "0.1.1"
@@ -8267,15 +8456,15 @@ dependencies = [

 [[package]]
 name = "png"
-version = "0.17.10"
+version = "0.17.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dd75bf2d8dd3702b9707cdbc56a5b9ef42cec752eb8b3bafc01234558442aa64"
+checksum = "82151a2fc869e011c153adc57cf2789ccb8d9906ce52c0b39a6b5697749d7526"
 dependencies = [
 "bitflags 1.3.2",
 "crc32fast",
 "fdeflate",
 "flate2",
- "miniz_oxide 0.7.2",
+ "miniz_oxide 0.8.4",
 ]

 [[package]]
@@ -8360,6 +8549,7 @@ dependencies = [
 name = "primitives"
 version = "0.0.1"
 dependencies = [
+ "assertor",
 "my-workspace-hack",
 "testing",
 ]
@@ -10382,6 +10572,17 @@ version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e"

+[[package]]
+name = "simplelog"
+version = "0.12.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "16257adbfaef1ee58b1363bdc0664c9b8e1e30aed86049635fb5f147d065a9c0"
+dependencies = [
+ "log",
+ "termcolor",
+ "time",
+]
+
 [[package]]
 name = "siphasher"
 version = "0.3.11"
@@ -10471,8 +10672,8 @@ dependencies = [
 "foreign-types 0.5.0",
 "js-sys",
 "log",
- "objc2",
- "objc2-foundation",
+ "objc2 0.5.2",
+ "objc2-foundation 0.2.2",
 "objc2-quartz-core",
 "raw-window-handle",
 "redox_syscall 0.5.2",
@@ -11644,9 +11845,9 @@ dependencies = [
 "log",
 "mime",
 "muda",
- "objc2",
+ "objc2 0.5.2",
 "objc2-app-kit",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 "percent-encoding",
 "plist",
 "raw-window-handle",
@@ -11754,16 +11955,16 @@ dependencies = [

 [[package]]
 name = "tauri-plugin-log"
-version = "2.2.1"
+version = "2.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "367a28a5e0ca39eac98005699466e8906edc4a2a8f8e13a5f1a71dc0bea6c677"
+checksum = "2341d5b9bc5318c8e34f35a569140c78337241aa9c14091550b424c49f0314e0"
 dependencies = [
 "android_logger",
 "byte-unit",
 "fern",
 "log",
- "objc2",
- "objc2-foundation",
+ "objc2 0.6.0",
+ "objc2-foundation 0.3.0",
 "serde",
 "serde_json",
 "serde_repr",
@@ -11803,9 +12004,9 @@ dependencies = [
 "http 1.1.0",
 "jni",
 "log",
- "objc2",
+ "objc2 0.5.2",
 "objc2-app-kit",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 "percent-encoding",
 "raw-window-handle",
 "softbuffer",
@@ -11912,6 +12113,15 @@ dependencies = [
 "safetensors",
 ]

+[[package]]
+name = "termcolor"
+version = "1.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755"
+dependencies = [
+ "winapi-util",
+]
+
 [[package]]
 name = "termtree"
 version = "0.4.1"
@@ -12573,9 +12783,9 @@ dependencies = [

 [[package]]
 name = "tracing-subscriber"
-version = "0.3.19"
+version = "0.3.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e8189decb5ac0fa7bc8b96b7cb9b2701d60d48805aca84a238004d665fcc4008"
+checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b"
 dependencies = [
 "nu-ansi-term",
 "sharded-slab",
@@ -12606,9 +12816,9 @@ dependencies = [
 "dirs 5.0.1",
 "libappindicator",
 "muda",
- "objc2",
+ "objc2 0.5.2",
 "objc2-app-kit",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 "once_cell",
 "png",
 "serde",
@@ -12884,6 +13094,12 @@ version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd"

+[[package]]
+name = "unicode-xid"
+version = "0.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
+
 [[package]]
 name = "unicode_categories"
 version = "0.1.1"
@@ -13573,9 +13789,9 @@ version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3ea403deff7b51fff19e261330f71608ff2cdef5721d72b64180bb95be7c4150"
 dependencies = [
- "objc2",
+ "objc2 0.5.2",
 "objc2-app-kit",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 "raw-window-handle",
 "windows-sys 0.59.0",
 "windows-version",
@@ -13674,6 +13890,12 @@ dependencies = [
 "syn 2.0.87",
 ]

+[[package]]
+name = "windows-link"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "76840935b766e1b0a05c0066835fb9ec80071d4c09a16f6bd5f7e655e3c14c38"
+
 [[package]]
 name = "windows-registry"
 version = "0.2.0"
@@ -14068,9 +14290,9 @@ dependencies = [
 "kuchikiki",
 "libc",
 "ndk 0.9.0",
- "objc2",
+ "objc2 0.5.2",
 "objc2-app-kit",
- "objc2-foundation",
+ "objc2-foundation 0.2.2",
 "objc2-ui-kit",
 "objc2-web-kit",
 "once_cell",
@@ -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",
@@ -56,6 +58,7 @@ members = [
  "crates/lib/files/mimetypes",
  "crates/lib/files/videos",
  "crates/lib/idempotency",
+  "crates/lib/image_algorithms",
  "crates/lib/jobs_common",
  "crates/lib/logging",
  "crates/lib/markdown",
@@ -0,0 +1,80 @@
+[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" }
+image_algorithms = { path = "../../lib/image_algorithms" }
+memory_caching = { path = "../../lib/caching/memory_caching" }
+ml_weights_registry = { path = "../ml_weights_registry" }
+primitives = { path = "../../lib/data/primitives" }
+
+# 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"
@@ -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,66 @@
+use crate::ml::background_removal::onnx_session::onnx_session;
+use image::{imageops, DynamicImage};
+use ndarray::{Array, CowArray};
+use ort::Value;
+use std::path::Path;
+use image_algorithms::remove_image_border::remove_image_border;
+
+pub async fn remove_image_background<P: AsRef<Path>>(model_path: P, image: DynamicImage) -> anyhow::Result<DynamicImage> {
+  // NB: Real time needed cuDNN 9, but this uses cuDNN 8:
+  // 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)?;
+  
+  // Crop transparent border cruft away.
+  let image = remove_image_border(&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,297 @@
+use std::path::Path;
+use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
+
+use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
+use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
+use crate::ml::model_cache::ModelCache;
+use crate::ml::prompt_cache::PromptCache;
+use candle_nn::VarBuilder;
+use crate::ml::model_config::ModelConfig;
+use anyhow::{anyhow, Error as E, Result};
+use candle_core::{DType, IndexOp, Tensor, D, Device, Module};
+use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
+use candle_transformers::models::stable_diffusion::lcm::LCMScheduler;
+use image::DynamicImage;
+use log::{debug, info};
+use rand::Rng;
+use tauri::{AppHandle, Emitter};
+use candle_transformers::models::{clip, t5, flux};
+use tokenizers::Tokenizer;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum FluxModel {
+  Schnell,
+  Dev,
+}
+
+pub struct FluxPipelineArgs<'a> {
+  pub image: &'a DynamicImage,
+  pub prompt: String,
+  pub model_cache: &'a ModelCache,
+  pub prompt_cache: &'a PromptCache,
+  pub app: &'a AppHandle,
+  pub flux_model: FluxModel,
+  pub active_blocks: Option<usize>,
+  pub use_cuda_stream: Option<bool>,
+  pub prefetch_next_batch: Option<bool>,
+  pub use_full_gpu: Option<bool>,
+  pub img2img_strength: Option<f64>,
+  pub seed: Option<u64>,
+  pub use_quantized_model: Option<bool>,
+  pub model_config: &'a ModelConfig,
+  pub flux_dev_or_schnell_path: &'a Path,
+  pub flux_schnell_quantized_gguf: &'a Path,
+  pub flux_autoencoder_path: &'a Path,
+  pub google_t5_model_path: &'a Path,
+  pub google_t5_config_path: &'a Path,
+  pub lmz_t5_tokenizer_path: &'a Path,
+  pub openai_clip_model_path: &'a Path,
+  pub openai_clip_tokenizer_json_path: &'a Path,
+}
+
+pub fn flux_pipeline(args: FluxPipelineArgs<'_>) -> Result<RgbImage> {
+  let FluxPipelineArgs { 
+    prompt, 
+    model_cache, 
+    prompt_cache, 
+    app, 
+    image, 
+    flux_model, 
+    active_blocks, 
+    use_cuda_stream, 
+    prefetch_next_batch, 
+    use_full_gpu, 
+    img2img_strength, 
+    seed, 
+    use_quantized_model, 
+    model_config ,
+    flux_dev_or_schnell_path,
+    flux_schnell_quantized_gguf,
+    flux_autoencoder_path,
+    google_t5_model_path,
+    google_t5_config_path,
+    lmz_t5_tokenizer_path,
+    openai_clip_model_path,
+    openai_clip_tokenizer_json_path,
+  } = args;
+
+  if use_quantized_model.unwrap_or(false) {
+    #[cfg(feature = "cuda")]
+    candle_core::quantized::cuda::set_force_dmmv(false);
+  }
+
+  let height = 1024;
+  let width = 1024;
+
+  let quantized = use_quantized_model.unwrap_or(false);
+  let device = model_config.device.clone();
+  let dtype = device.bf16_default_to_f32();
+
+  info!("Starting image generation with FLUX model:");
+  info!("  Model: {:?}", flux_model);
+  info!("  Prompt: {}", prompt);
+  info!("  Dimensions: {}x{}", width, height);
+  info!("  Device: {:?}", model_config.device);
+
+  let seed_value = seed.unwrap_or_else(|| rand::thread_rng().gen());
+  device.set_seed(seed_value)?;
+  
+  info!("Using seed: {}", seed_value);
+
+  // Setup FLUX model configuration
+  let flux_config = match flux_model {
+    FluxModel::Dev => flux::model::Config::dev(),
+    FluxModel::Schnell => flux::model::Config::schnell(),
+  };
+
+  // Get T5 embeddings
+  let t5_emb = {
+    info!("Generating T5 embeddings for prompt...");
+    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[google_t5_model_path], dtype , &device)? };
+    let config = std::fs::read_to_string(google_t5_config_path)?;
+    let config: t5::Config = serde_json::from_str(&config)?;
+
+    let mut model = t5::T5EncoderModel::load(vb, &config)?;
+    let tokenizer = Tokenizer::from_file(lmz_t5_tokenizer_path).map_err(E::msg)?;
+
+    let mut tokens = tokenizer.encode(prompt.as_str(), true).map_err(E::msg)?.get_ids().to_vec();
+    tokens.resize(256, 0);
+    let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
+
+    model.forward(&input_token_ids)?
+  };
+  
+  info!("T5 embeddings generated");
+
+  // Get CLIP embeddings
+  let clip_emb = {
+    info!("Generating CLIP embeddings for prompt...");
+    let vb =
+        unsafe { VarBuilder::from_mmaped_safetensors(&[openai_clip_model_path], dtype, &device)? };
+    // https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json
+    let config = clip::text_model::ClipTextConfig {
+        vocab_size: 49408,
+        projection_dim: 768,
+        activation: clip::text_model::Activation::QuickGelu,
+        intermediate_size: 3072,
+        embed_dim: 768,
+        max_position_embeddings: 77,
+        pad_with: None,
+        num_hidden_layers: 12,
+        num_attention_heads: 12,
+    };
+    let model =
+        clip::text_model::ClipTextTransformer::new(vb.pp("text_model"), &config)?;
+    let tokenizer = Tokenizer::from_file(openai_clip_tokenizer_json_path).map_err(E::msg)?;
+    let tokens = tokenizer
+        .encode(prompt.as_str(), true)
+        .map_err(E::msg)?
+        .get_ids()
+        .to_vec();
+    let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
+    model.forward(&input_token_ids)?
+  };
+  info!("CLIP embeddings generated");
+
+  // Process input image to tensor
+  info!("Converting input image to tensor...");
+  let img_tensor = dynamic_image_to_tensor(image, &device, dtype)?;
+  info!("Image tensor shape: {:?}", img_tensor.shape());
+
+  // Create initial noise - this is important to match the expected format
+  info!("Generating initial noise for FLUX...");
+  let initial_tensor = flux::sampling::get_noise(1, height, width, &device)?.to_dtype(dtype)?;
+  info!("Initial tensor shape: {:?}", initial_tensor.shape());
+  
+  // Blend initial noise with input image if img2img_strength is specified
+  let initial_tensor = if let Some(strength) = img2img_strength {
+    if strength > 0.0 && strength < 1.0 {
+      info!("Using img2img with strength: {}", strength);
+      // Reshape the input tensor to match the noise tensor dimensions
+      let adjusted_img = if img_tensor.shape() != initial_tensor.shape() {
+        info!("Reshaping input image to match noise dimensions...");
+        // You may need more complex reshaping logic here depending on your use case
+        let (_, c, h, w) = initial_tensor.dims4()?;
+        img_tensor.upsample_nearest2d(h, w)?
+      } else {
+        img_tensor
+      };
+      
+      // Blend the noise and the image according to strength
+      let strength_tensor = Tensor::new(strength as f32, &device)?;
+      let inverse_strength = Tensor::new(1.0 - strength as f32, &device)?;
+      let weighted_noise = (&initial_tensor * &strength_tensor)?;
+      let weighted_img = (&adjusted_img * &inverse_strength)?;
+      (&weighted_noise + &weighted_img)?
+    } else {
+      initial_tensor
+    }
+  } else {
+    initial_tensor
+  };
+  info!("Final initial tensor shape: {:?}", initial_tensor.shape());
+
+  // Create FLUX state
+  info!("Creating FLUX state...");
+  let state = if quantized {
+    flux::sampling::State::new(
+        &t5_emb.to_dtype(DType::F32)?,
+        &clip_emb.to_dtype(DType::F32)?,
+        &initial_tensor.to_dtype(DType::F32)?,
+    )?
+  } else {
+    flux::sampling::State::new(&t5_emb, &clip_emb, &initial_tensor)?
+  };
+  info!("{state:?}");
+
+  // Set appropriate timesteps based on model
+  let timesteps = match flux_model {
+    FluxModel::Dev => flux::sampling::get_schedule(28, Some((state.img.dim(1)?, 0.5, 1.15))),
+    FluxModel::Schnell => flux::sampling::get_schedule(4, None),
+  };
+
+  info!("Starting FLUX denoising with timesteps: {:?}", timesteps);
+
+  // Actual model loading and inference
+  let img_latent = if quantized {
+    // Quantized model path
+    info!("Loading quantized FLUX model...");
+    let model_file = match flux_model {
+      FluxModel::Schnell => flux_schnell_quantized_gguf,
+      FluxModel::Dev => return Err(anyhow!("Quantized model not available for FLUX.1-dev")),
+    };
+
+    info!("Loading quantized model from {}", model_file.to_str().unwrap());
+
+    let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(model_file, &device)?;
+    let mut model = flux::quantized_model::Flux::new(&flux_config, vb)?;
+    info!("Quantized model loaded on {:?}", device);
+
+    // Run denoising
+    info!("Running denoising...");
+    let img_latent = flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, &timesteps, 4.)?.to_dtype(dtype)?;
+    info!("Denoising completed");
+    img_latent
+  } else {
+    // Full precision model
+    info!("Loading FLUX model...");
+
+    let mut model = if use_full_gpu.unwrap_or(false) {
+      // Load entire model on GPU
+      info!("Loading entire model directly to GPU...");
+      let gpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[flux_dev_or_schnell_path], dtype, &device)? };
+      let mut model = flux::model::Flux::new(&flux_config, gpu_vb)?;
+      model.set_use_device_management(false);
+      model
+    } else {
+      // Use memory-efficient approach with blocks on CPU
+      info!("Using memory-efficient approach with blocks on CPU...");
+      let cpu_device = candle_core::Device::Cpu;
+      let cpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[flux_dev_or_schnell_path], dtype, &cpu_device)? };
+      let mut model = flux::model::Flux::new_with_gpu_core(&flux_config, cpu_vb, &device)?;
+
+      // Set active blocks
+      let active_blocks_count = active_blocks.unwrap_or(1);
+      model.set_max_active_blocks(active_blocks_count);
+      info!("Using {} active blocks on GPU", active_blocks_count);
+
+      // Enable prefetching if requested
+      if prefetch_next_batch.unwrap_or(false) {
+        model.set_prefetch_next_batch(true);
+        info!("Prefetching enabled - will overlap transfers with computation");
+      }
+
+      model
+    };
+
+    // Run denoising
+    flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, &timesteps, 4.)?
+  };
+
+  info!("FLUX denoising completed, unpacking latent image...");
+  let img_unpacked = flux::sampling::unpack(&img_latent, height, width)?;
+
+  // Load autoencoder for final decoding
+  info!("Loading autoencoder for final decoding...");
+  
+  let ae_config = match flux_model {
+    FluxModel::Dev => flux::autoencoder::Config::dev(),
+    FluxModel::Schnell => flux::autoencoder::Config::schnell(),
+  };
+
+  let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[flux_autoencoder_path], dtype, &device)? };
+  let ae_model = flux::autoencoder::AutoEncoder::new(&ae_config, vb)?;
+
+  info!("Decoding final image...");
+  let img = ae_model.decode(&img_unpacked)?;
+
+  // Convert tensor to image
+  info!("Converting tensor to image...");
+  let img = ((img.clamp(-1f32, 1f32)? + 1.0)? * 127.5)?.to_dtype(candle_core::DType::U8)?;
+
+  // Convert to RGB image
+  let image = tensor_to_image_buffer(&img.i(0)?)?;
+
+  info!("FLUX pipeline completed successfully");
+  Ok(image)
+}
@@ -0,0 +1 @@
+pub mod flux_pipeline;
@@ -0,0 +1,8 @@
+pub mod background_removal;
+pub mod flux;
+pub mod image;
+pub mod model_cache;
+pub mod model_config;
+pub mod model_file;
+pub mod prompt_cache;
+pub mod stable_diffusion;
@@ -1,15 +1,14 @@
-use crate::ml::models::unet_model::UNetModel;
 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<UNetModel>>>,
+  unet2: RwLock<Option<Arc<UNet2DConditionModel>>>,
 }

 impl ModelCache {
@@ -39,7 +38,7 @@ impl ModelCache {
    }
  }

-  pub fn set_unet(&self, vae: Arc<UNetModel>) -> anyhow::Result<()> {
+  pub fn set_unet(&self, vae: Arc<UNet2DConditionModel>) -> anyhow::Result<()> {
    match self.unet2.write() {
      Err(err) => Err(anyhow!("{}", err)),
      Ok(mut lock) => {
@@ -49,7 +48,7 @@ impl ModelCache {
    }
  }

-  pub fn get_unet(&self) -> anyhow::Result<Option<Arc<UNetModel>>> {
+  pub fn get_unet(&self) -> anyhow::Result<Option<Arc<UNet2DConditionModel>>> {
    match self.unet2.read() {
      Err(err) => Err(anyhow!("{}", err)),
      Ok(ref lock) => {
@@ -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,
+    })
+  }
+}
@@ -3,19 +3,51 @@ extern crate accelerate_src;
 #[cfg(feature = "mkl")]
 extern crate intel_mkl_src;

-use std::path::PathBuf;
 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;
-use crate::ml::model_type::ModelType;
-use crate::state::app_dir::AppWeightsDir;
+
+
+pub struct ClipArgs<'a, P: AsRef<Path>,  Q: AsRef<Path>> {
+  pub prompt: &'a str, 
+  pub uncond_prompt: &'a str, 
+  pub tokenizer: Option<String>, 
+  pub clip_weights: Option<String>, 
+  pub clip2_weights: Option<String>, 
+  pub sd_version: StableDiffusionVersion, 
+  pub sd_config: &'a stable_diffusion::StableDiffusionConfig, 
+  pub use_f16: bool, 
+  pub device: &'a Device, 
+  pub dtype: DType, 
+  pub use_guide_scale: bool, 
+  pub clip_json_path: P,
+  pub clip_weights_path: Q,
+}

 #[allow(clippy::too_many_arguments)]
-pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: Option<String>, clip_weights: Option<String>, clip2_weights: Option<String>, sd_version: StableDiffusionVersion, sd_config: &stable_diffusion::StableDiffusionConfig, use_f16: bool, device: &Device, dtype: DType, use_guide_scale: bool, weights_dir: &AppWeightsDir) -> anyhow::Result<Tensor> {
+pub fn infer_clip_text_embeddings<P: AsRef<Path>, Q: AsRef<Path>>(args: ClipArgs<'_, P, Q>) -> anyhow::Result<Tensor> {
+  let ClipArgs {
+    prompt, 
+    uncond_prompt, 
+    tokenizer, 
+    clip_weights, 
+    clip2_weights, 
+    sd_version, 
+    sd_config, 
+    use_f16, 
+    device, 
+    dtype, 
+    use_guide_scale, 
+    clip_json_path,
+    clip_weights_path,
+  } = args;
+
  info!("Preparing text embeddings... for {:?}", sd_version);

  let which = match sd_version {
@@ -39,7 +71,8 @@ pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer:
        dtype,
        use_guide_scale,
        *first,
-        weights_dir,
+        clip_json_path.as_ref(),
+        clip_weights_path.as_ref(),
      )
    })
    .collect::<anyhow::Result<Vec<_>>>()?;
@@ -50,20 +83,25 @@ pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer:
 }

 #[allow(clippy::too_many_arguments)]
-fn do_infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: Option<String>, clip_weights: Option<String>, clip2_weights: Option<String>, sd_version: StableDiffusionVersion, sd_config: &stable_diffusion::StableDiffusionConfig, use_f16: bool, device: &Device, dtype: DType, use_guide_scale: bool, first: bool, weights_dir: &AppWeightsDir) -> anyhow::Result<Tensor> {
+fn do_infer_clip_text_embeddings(
+  prompt: &str, 
+  uncond_prompt: &str, 
+  tokenizer: Option<String>, 
+  clip_weights: Option<String>, 
+  clip2_weights: Option<String>, 
+  sd_version: StableDiffusionVersion, 
+  sd_config: &stable_diffusion::StableDiffusionConfig, 
+  use_f16: bool, 
+  device: &Device, 
+  dtype: DType, 
+  use_guide_scale: bool, 
+  first: bool, 
+  clip_json_path: &Path,
+  clip_weights_path: &Path,
+) -> anyhow::Result<Tensor> {
  info!("do_infer_clip_text_embeddings called with args {:?}", (prompt, uncond_prompt, tokenizer.clone().unwrap_or_else(|| "None".to_string()), clip_weights.clone().unwrap_or_else(|| "None".to_string()), clip2_weights.clone().unwrap_or_else(|| "None".to_string()), sd_version, sd_config, use_f16, device, dtype, use_guide_scale, first));

-  //let tokenizer_file = if first {
-  //  info!("ModelFile::Tokenizer");
-  //  ModelFile::Tokenizer
-  //} else {
-  //  info!("ModelFile::Tokenizer2");
-  //  ModelFile::Tokenizer2
-  //};
-
-  //let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
-
-  let tokenizer = weights_dir.model_path(&ModelType::ClipJson);
+  let tokenizer = clip_json_path;

  info!("Loading Clip Tokenizer path: {:?}", tokenizer);

@@ -85,14 +123,16 @@ fn do_infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: O
  let tokens = Tensor::new(tokens.as_slice(), device)?.unsqueeze(0)?;

  println!("Building the Clip transformer.");
-  let clip_weights_file = if first { ModelFile::Clip } else { ModelFile::Clip2 };
+  
+  // TODO(bt,2025-03-23): These used to be two different clip models. We should only need one.
  let clip_weights = if first {
-    //clip_weights_file.get(clip_weights, sd_version, use_f16)?
-    weights_dir.model_path(&ModelType::SdxlTurboClipEncoder)
+    println!(">>>>> CLIP 1");
+    clip_weights_path
  } else {
-    //clip_weights_file.get(clip2_weights, sd_version, use_f16)?
-    weights_dir.model_path(&ModelType::SdxlTurboClipEncoder2)
+    println!(">>>>> CLIP 2");
+    clip_weights_path
  };
+  
  let clip_config = if first { &sd_config.clip } else { sd_config.clip2.as_ref().unwrap() };
  let text_model = stable_diffusion::build_clip_transformer(clip_config, clip_weights, device, DType::F32)?;

@@ -1,72 +1,100 @@
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};

-use crate::events::notification_event::{NotificationEvent, NotificationModelType};
-use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
-use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
-use crate::ml::model_cache::ModelCache;
-use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
-use crate::ml::model_type::ModelType;
-use crate::ml::models::unet_model::UNetModel;
-use crate::ml::prompt_cache::PromptCache;
-use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
-use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_embeddings;
-use crate::state::app_config::AppConfig;
-use crate::state::app_dir::AppDataRoot;
 use anyhow::{anyhow, Error as E, Result};
-use candle_core::{DType, IndexOp, Tensor, D};
+use candle_core::{DType, Device, IndexOp, Tensor, D};
 use candle_nn::VarBuilder;
-use candle_transformers::models::stable_diffusion::lcm::LCMScheduler;
+use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
+use candle_transformers::models::stable_diffusion::lcm::{LCMScheduler, LCMSchedulerConfig};
 use candle_transformers::models::stable_diffusion::unet_2d::BlockConfig;
 use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
 use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModelConfig;
 use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
+use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
+use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
+use crate::ml::model_cache::ModelCache;
+use crate::ml::model_config::ModelConfig;
+use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
+use crate::ml::prompt_cache::PromptCache;
+use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
+use crate::ml::stable_diffusion::infer_clip_text_embeddings::{infer_clip_text_embeddings, ClipArgs};
+use crate::ml::stable_diffusion::remap_lcm_strength_range::remap_lcm_strength_range;
 use image::DynamicImage;
-use log::info;
+use log::{debug, error, info};
+use ml_weights_registry::weights_registry::weights::{LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
 use rand::Rng;
-use tauri::{AppHandle, Emitter};

-pub struct Args<'a> {
+const DEFAULT_STRENGTH : f64 = 75.0;
+
+const DEFAULT_SEED : u64 = 42;
+
+pub struct Args<'a, P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>> {
  pub image: &'a DynamicImage,
  pub prompt: String,
  pub uncond_prompt: String,
  pub cfg_scale: Option<f64>,
-  pub i2i_strength: Option<u8>,
-  pub configs: &'a AppConfig,
+  pub img2img_strength: Option<f64>,
  pub model_cache: &'a ModelCache,
  pub prompt_cache: &'a PromptCache,
-  pub app_data_root: &'a AppDataRoot,
-  pub app: &'a AppHandle,
  pub use_flash_attn: bool,
+  pub model_config: &'a ModelConfig,
+  pub maybe_seed: Option<u64>,
+  pub scheduler_steps: usize,
+  pub vae_path: P1,
+  pub unet_path: P2,
+  pub clip_json_path: P3,
+  pub clip_weights_path: P4,
 }

-pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
-  let Args { prompt, uncond_prompt, cfg_scale, i2i_strength, configs, model_cache, prompt_cache, app, image, app_data_root, use_flash_attn } = args;
+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 weights_dir = app_data_root.weights_dir();
+  let img2img_strength = img2img_strength
+    .map(|s| remap_lcm_strength_range(s))
+    .unwrap_or(DEFAULT_STRENGTH);
+  
+  let img2img_strength = img2img_strength / 100.0;

-  let img2img_strength = args.i2i_strength.unwrap_or(75) as f64 / 100.0;
  let use_f16 = true;

-  //println!("Starting image generation with the following configuration:");
-  //println!("  Model: {:?}", configs.sd_version);
-  //println!("  Prompt: {}", prompt);
-  //println!("  Steps: {}", configs.scheduler_steps);
-  //println!("  Device: {:?}", configs.device);
-
-  //println!("Model dimensions: {}x{}", configs.sd_config.width, configs.sd_config.height);
-
  // Use LCM Scheduler instead of Euler Ancestral for better speed and quality
-  let mut scheduler = LCMScheduler::new(configs.scheduler_steps, img2img_strength, candle_transformers::models::stable_diffusion::lcm::LCMSchedulerConfig::default())?;
+  let mut scheduler = LCMScheduler::new(scheduler_steps, img2img_strength, LCMSchedulerConfig::default())?;

-  let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
-  configs.device.set_seed(seed)?;
+  //let seed = maybe_seed.unwrap_or_else(|| rand::thread_rng().gen());
+  let seed = maybe_seed.unwrap_or(DEFAULT_SEED);

-  //info!("Using seed: {}", seed);
+  device.set_seed(seed)?;

  let guidance_scale = match cfg_scale {
    Some(guidance_scale) => guidance_scale,
-    None => match configs.sd_version {
+    None => match sd_version {
      StableDiffusionVersion::V1_5 | StableDiffusionVersion::V2_1 | StableDiffusionVersion::Xl => 7.5,
      StableDiffusionVersion::Turbo => 0.,
      _ => 0., // NB: Not sure what the other model families should use, so sticking with "0"
@@ -76,72 +104,50 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
  // let use_guide_scale = guidance_scale > 1.0;
  let use_guide_scale = false;

-  //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!("Loading clip weights...");
-    //println!("Loading clip weights for {:?} with model {:?}", configs.sd_version, configs.sd_version.repo());
-    //let api = configs.hf_api.clone();
-
-    //let clip_path = api.model(configs.sd_version.repo().to_string()).get(if use_f16 { "text_encoder/model.fp16.safetensors" } else { "text_encoder/model.safetensors" })?;
-
-    //println!(">>>> CLIP REPO = {:?}", configs.sd_version.repo());
-    //println!(">>>> CLIP PATH = {:?}", clip_path);
-    //println!(">>>> USE_FP16 = {:?}", use_f16);
-    //let clip_weights = Some(clip_path.to_string_lossy().to_string());
-
-    let clip_weights = weights_dir.model_path(&ModelType::LykonDreamshaper7TextEncoderFp16);
-    let clip_weights = Some(clip_weights.to_string_lossy().to_string());
-
    info!("Prompt is NOT cached! Calculating embedding...");
-    let tensor = infer_clip_text_embeddings(
-      &prompt,
-      &uncond_prompt,
-      None, // tokenizer
-      clip_weights,
-      None, // clip2_weights
-      configs.sd_version,
-      &configs.sd_config,
+    let tensor = infer_clip_text_embeddings(ClipArgs {
+      prompt: & prompt,
+      uncond_prompt: &uncond_prompt,
+      tokenizer: None, // tokenizer
+      clip_weights: None,
+      clip2_weights: None, // clip2_weights
+      sd_version,
+      sd_config: &sd_config,
      use_f16, // use_f16
-      &configs.device,
-      configs.dtype,
+      device,
+      dtype,
      use_guide_scale,
-      weights_dir,
-    )?;
+      clip_json_path,
+      clip_weights_path,
+    })?;
+
    prompt_cache.store_copy(&prompt, &tensor)?;
    tensor
  };

-  println!("Text embeddings shape: {:?}", text_embeddings.shape());
+  debug!("Text embeddings shape: {:?}", text_embeddings.shape());

-  println!("Loading input image into tensor...");
+  debug!("Loading input image into tensor...");

-  let input_image = dynamic_image_to_tensor(image, &configs.device, configs.dtype)?;
+  let input_image = dynamic_image_to_tensor(image, &device, dtype)?;

-  println!("Reference image shape: {:?}", input_image.shape());
+  debug!("Reference image shape: {:?}", input_image.shape());

-  let vae_scale = get_vae_scale(configs.sd_version);
+  let vae_scale = get_vae_scale(sd_version);

  let maybe_vae = model_cache.get_vae()?;

  let vae = match maybe_vae {
    Some(vae) => vae,
    None => {
-      //let repo = configs.sd_version.repo();
-      //println!(">>>>> VAE REPO = {:?}", repo);
-      //println!("Building VAE model from : {:?} ... (3)", repo);
-      //let vae_file = configs.hf_api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?;
-
-      let vae_file = weights_dir.model_path(&ModelType::LykonDreamshaper7Vae);
-
-      println!("Building VAE model from file {:?}...", &vae_file);
-
-      let vae = configs.sd_config.build_vae(vae_file, &configs.device, configs.dtype)?;
+      info!("Building VAE model from file {:?}...", vae_path.as_ref());
+      
+      let vae = sd_config.build_vae(vae_path, &device, dtype)?;

      let vae = Arc::new(vae);

@@ -156,11 +162,9 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
  let unet = match maybe_unet {
    Some(unet) => unet,
    None => {
-      info!("No unet found in cache; loading...");
+      info!("Building UNET model from file {:?}...", unet_path.as_ref());

-      let unet_weights = weights_dir.model_path(&ModelType::SimianLuoLcmDreamshaperV7Unet);
-
-      let in_channels = match configs.sd_version {
+      let in_channels = match sd_version {
        StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::V1_5Inpaint => 9,
        _ => 4,
      };
@@ -195,12 +199,10 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
      };

      // Load the model directly from the safetensors file
-      let vs_unet = unsafe { VarBuilder::from_mmaped_safetensors(&[unet_weights], configs.dtype, &configs.device)? };
+      let vs_unet = unsafe { VarBuilder::from_mmaped_safetensors(&[unet_path], dtype, &device)? };

      let unet = UNet2DConditionModel::new(vs_unet, in_channels, 4, use_flash_attn, unet_config)?;

-      let unet = UNetModel::new_with_inner(unet);
-
      let unet = Arc::new(unet);

      model_cache.set_unet(unet.clone())?;
@@ -211,50 +213,40 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {

  let init_latent_dist: DiagonalGaussianDistribution = vae.encode(&input_image)?;

-  // Parse and calculate the img2img strength
-  let img2img_strength = match i2i_strength {
-    None => 0.75f64,
-    Some(strength) => (strength as f64) / 100.0f64,
-  };
-
  // Create LCM guidance scale embeddings
  let embedding_dim = 1280;
-  let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &configs.device, configs.dtype)?;
+  let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &device, dtype)?;

  // Generate latents from input image using the LCM approach
-  println!("Generating latents from input image...");
+  debug!("Generating latents from input image...");
  let init_latents = (init_latent_dist.sample()? * vae_scale)?;
-  println!("Initial latents shape: {:?}", init_latents.shape());
+
+  debug!("Initial latents shape: {:?}", init_latents.shape());

  // Calculate timesteps for LCM with proper img2img handling
-  let timesteps = LCMScheduler::get_timesteps_for_steps(configs.scheduler_steps, img2img_strength);
+  let timesteps = LCMScheduler::get_timesteps_for_steps(scheduler_steps, img2img_strength);
  let t_start = timesteps[0]; // First denoising step

  // Add noise at the appropriate timestep for img2img
-  println!("Adding noise to latents for t_start={}", t_start);
+  debug!("Adding noise to latents for t_start={}", t_start);
  let noise = init_latents.randn_like(0f64, 1f64)?;
  let latents = scheduler.add_noise(&init_latents, noise, t_start)?;

-  println!("Latents initialized successfully");
-  println!("Initial noise shape: {:?}", latents.shape());
-
-  println!("Starting diffusion process...");
-
  let mut latents = latents;
  let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect();

  for (timestep_index, &timestep) in timesteps.iter().enumerate() {
-    println!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep);
+    info!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep);

    let latent_model_input = if use_guide_scale { Tensor::cat(&[&latents, &latents], 0)? } else { latents.clone() };

    let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep);

    // Use the guidance scale embedding in the UNet inference with the proper method name
-    let mut noise_pred = match unet.forward_with_guidance(&latent_model_input, timestep as f64, &text_embeddings, &guidance_scale_embedding) {
+    let mut noise_pred = match unet.forward_with_guidance(&latent_model_input, timestep as f64, &text_embeddings, Some(&guidance_scale_embedding)) {
      Ok(pred) => pred,
      Err(e) => {
-        println!("UNet inference failed with error: {}", e);
+        error!("UNet inference failed with error: {}", e);
        return Err(anyhow::anyhow!("UNet inference failed: {}", e));
      },
    };
@@ -266,24 +258,22 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
    }

    // Apply the LCM scheduler step
-    latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, configs.scheduler_steps)?;
+    latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, scheduler_steps)?;
  }

-  println!("Diffusion process completed, decoding image...");
+  debug!("Diffusion process completed, decoding image...");
  let image = vae.decode(&(latents / vae_scale)?)?;

-  println!("VAE decode completed");
-
-  println!("Scaling image back");
+  debug!("Scaling image back");
  let image = ((image / 2.)? + 0.5)?;

-  println!("Normalized image values");
+  debug!("Normalized image values");
  let image = (image.clamp(0f32, 1.)? * 255.)?;

-  println!("Convert to int8");
+  debug!("Convert to int8");
  let image = image.to_dtype(DType::U8)?;

-  println!("Converted to 8-bit format");
+  debug!("Converted to 8-bit format");

  let image = tensor_to_image_buffer(&image.i(0)?)?;

@@ -1,4 +1,5 @@
 pub mod get_vae_scale;
 pub mod infer_clip_text_embeddings;
 pub mod stable_diffusion_pipeline;
-pub mod lcm_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);
+  }
+}
@@ -1,76 +1,78 @@
+use std::path::Path;
 use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};

-use crate::events::notification_event::{NotificationModelType, NotificationEvent};
 use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
 use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
 use crate::ml::model_cache::ModelCache;
 use crate::ml::model_file::StableDiffusionVersion;
-use crate::ml::model_type::ModelType;
-use crate::ml::models::unet_model::UNetModel;
 use crate::ml::prompt_cache::PromptCache;
 use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
-use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_embeddings;
-use crate::state::app_config::AppConfig;
+use crate::ml::stable_diffusion::infer_clip_text_embeddings::{infer_clip_text_embeddings, ClipArgs};
 use anyhow::{anyhow, Error as E, Result};
-use candle_core::{DType, IndexOp, Tensor, D};
+use candle_core::{DType, Device, IndexOp, Tensor, D};
 use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
+use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
 use image::DynamicImage;
 use log::info;
 use rand::Rng;
-use tauri::{AppHandle, Emitter};
-use crate::state::app_dir::AppDataRoot;

-pub struct Args<'a> {
+pub struct Args<'a, P: AsRef<Path>, Q: AsRef<Path>, R: AsRef<Path>, S: AsRef<Path>> {
    pub image: &'a DynamicImage,
    pub prompt: String,
    pub uncond_prompt: String,
    pub cfg_scale: Option<f64>,
    pub i2i_strength: Option<u8>,
-    pub configs: &'a AppConfig,
    pub model_cache: &'a ModelCache,
    pub prompt_cache: &'a PromptCache,
-    pub app_data_root: &'a AppDataRoot,
-    pub app: &'a AppHandle,
+    pub sd_version: StableDiffusionVersion,
+    pub sd_config: StableDiffusionConfig,
+    pub device: &'a Device,
+    pub dtype: DType,
+    pub maybe_seed: Option<u64>,
+    pub scheduler_steps: usize,
+    pub sdxl_turbo_vae_path: P,
+    pub sdxl_turbo_unet_path: Q,
+    pub clip_json_path: R,
+    pub clip_weights_path: S,
 }

-pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
+pub fn stable_diffusion_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<RgbImage> 
+  where P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>
+{
    let Args { 
        prompt, 
        uncond_prompt, 
        cfg_scale, 
        i2i_strength, 
-        configs, 
        model_cache, 
        prompt_cache, 
-        app, 
        image,
-        app_data_root,
+        sd_version,
+        sd_config,
+        device,
+        dtype,
+        maybe_seed,
+        scheduler_steps,
+        sdxl_turbo_vae_path,
+        sdxl_turbo_unet_path,
+        clip_json_path,
+        clip_weights_path,
    } = args;
    
-    let weights_dir = app_data_root.weights_dir();
-    
-    println!("Starting image generation with the following configuration:");
-    println!("  Model: {:?}", configs.sd_version);
-    println!("  Prompt: {}", prompt);
-    println!("  Steps: {}", configs.scheduler_steps);
-    println!("  Device: {:?}", configs.device);
-
-    println!("Model dimensions: {}x{}", configs.sd_config.width, configs.sd_config.height);
-
    // TODO(bt,2025-02-18): The scheduler is `EulerAncestralDiscreteScheduler`, but we may want to port an LCM scheduler.
    //  This is a target for performance improvement
    //  See: https://github.com/huggingface/candle/issues/1331
-    let mut scheduler = configs.sd_config.build_scheduler(
-        configs.scheduler_steps)?;
+    let mut scheduler = sd_config.build_scheduler(scheduler_steps)?;

-    let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
-    configs.device.set_seed(seed)?;
+    let seed = maybe_seed.unwrap_or_else(|| rand::thread_rng().gen());
+    
+    device.set_seed(seed)?;

    info!("Using seed: {}", seed);

    let guidance_scale = match cfg_scale {
        Some(guidance_scale) => guidance_scale,
-        None => match configs.sd_version {
+        None => match sd_version {
            StableDiffusionVersion::V1_5
            | StableDiffusionVersion::V2_1
            | StableDiffusionVersion::Xl => 7.5,
@@ -91,20 +93,21 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
    } else {
        
        info!("Prompt is NOT cached! Calculating embedding...");
-        let tensor = infer_clip_text_embeddings(
-            &prompt,
-            &uncond_prompt,
-            None, // tokenizer
-            None, // clip_weights
-            None, // clip2_weights
-            configs.sd_version,
-            &configs.sd_config,
-            false, // use_f16
-            &configs.device,
-            configs.dtype,
+        let tensor = infer_clip_text_embeddings(ClipArgs {
+            prompt: &prompt,
+            uncond_prompt: &uncond_prompt,
+            tokenizer: None, // tokenizer
+            clip_weights: None, // clip_weights
+            clip2_weights: None, // clip2_weights
+            sd_version: sd_version,
+            sd_config: &sd_config,
+            use_f16: false, // use_f16
+            device: device,
+            dtype: dtype,
            use_guide_scale,
-            weights_dir,
-        )?;
+            clip_json_path: clip_json_path.as_ref(),
+            clip_weights_path: clip_weights_path.as_ref(),
+        })?;
        prompt_cache.store_copy(&prompt, &tensor)?;
        tensor
    };
@@ -113,14 +116,11 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {

    println!("Loading input image into tensor...");

-    let input_image = dynamic_image_to_tensor(
-        image,
-        &configs.device,
-        configs.dtype)?;
+    let input_image = dynamic_image_to_tensor(image, device, dtype)?;

    println!("Reference image shape: {:?}", input_image.shape());

-    let vae_scale = get_vae_scale(configs.sd_version);
+    let vae_scale = get_vae_scale(sd_version);

    let maybe_vae = model_cache.get_vae()?;
    
@@ -129,33 +129,14 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
        None => {
            info!("No vae found in cache; loading...");
            
-            let mut notify_download_complete = false;
-            //if !vae_file.exists() {
-            if true {
-                notify_download_complete = true;
-                app.emit("notification", NotificationEvent::ModelDownloadStarted {
-                    model_name: "sdxl-turbo-vae",
-                    model_type: NotificationModelType::Vae,
-                })?;
-            }
+            let vae_file = sdxl_turbo_vae_path.as_ref();
            
-            let vae_file = weights_dir.model_path(&ModelType::SdxlTurboVae);
-            
-            let vae = configs
-              .sd_config
-              .build_vae(vae_file, &configs.device, configs.dtype)?;
+            let vae = sd_config.build_vae(vae_file, &device, dtype)?;

            let vae = Arc::new(vae);
            
            model_cache.set_vae(vae.clone())?;
-            
-            if notify_download_complete {
-                app.emit("notification", NotificationEvent::ModelDownloadComplete {
-                    model_name: "sdxl-turbo-vae",
-                    model_type: NotificationModelType::Vae,
-                })?;
-            }
-            
+
            vae
        }
    };
@@ -167,32 +148,15 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
        None => {
            info!("No unet found in cache; loading...");

-            let mut notify_download_complete = false;
-            //if !unet_file.exists() {
-            if true {
-                notify_download_complete = true;
-                app.emit("notification", NotificationEvent::ModelDownloadStarted {
-                    model_name: "sdxl-turbo-unet",
-                    model_type: NotificationModelType::Unet,
-                })?;
-            }
-            
-            let unet_file = weights_dir.model_path(&ModelType::SdxlTurboUnet);
+            let unet_weights_file = sdxl_turbo_unet_path.as_ref();

-            let unet = UNetModel::new(&configs.sd_config, unet_file, &configs.device, configs.dtype)
+            let unet = sd_config.build_unet(unet_weights_file, &device, 4, true, dtype)
              .map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
-            
+
            let unet = Arc::new(unet);

            model_cache.set_unet(unet.clone())?;

-            if notify_download_complete {
-                app.emit("notification", NotificationEvent::ModelDownloadComplete {
-                    model_name: "sdxl-turbo-unet",
-                    model_type: NotificationModelType::Unet,
-                })?;
-            }
-
            unet
        }
    };
@@ -216,9 +180,9 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
    };

    let t_start = {
-        let start = configs.scheduler_steps - (configs.scheduler_steps as f64 * img2img_strength) as usize;
+        let start = scheduler_steps - (scheduler_steps as f64 * img2img_strength) as usize;

-        println!("Starting from step {} of {} (strength: {})", start, configs.scheduler_steps, img2img_strength);
+        println!("Starting from step {} of {} (strength: {})", start, scheduler_steps, img2img_strength);
        start
    };

@@ -257,7 +221,7 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {

        let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?;

-        let mut noise_pred = match unet.inference(&latent_model_input, timestep as f64, &text_embeddings) {
+        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);
@@ -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,80 @@
+use crate::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction};
+use crate::weights_registry::weight_descriptor_builder::weight;
+
+
+/// https://huggingface.co/black-forest-labs/FLUX.1-dev
+pub const FLUX_DEV : WeightDescriptor = weight!(
+  "Flux-1 Dev",
+  "flux1-dev.safetensors",
+  WeightFunction::ImageGeneration
+);
+
+// TODO(bt): I think the autoencoder for Dev and Schnell are the same.
+//  They have the same checksum unless I mixed up the files when downloading.
+
+/// https://huggingface.co/black-forest-labs/FLUX.1-dev
+pub const FLUX_DEV_AUTOENCODER : WeightDescriptor = weight!(
+  "Flux-1 Dev (Autoencoder)",
+  "flux1-dev.ae.safetensors",
+  WeightFunction::ImageGeneration
+);
+
+/// https://huggingface.co/black-forest-labs/FLUX.1-schnell
+pub const FLUX_SCHNELL : WeightDescriptor = weight!(
+  "Flux-1 Schnell",
+  "flux1-schnell.safetensors",
+  WeightFunction::ImageGeneration
+);
+
+/// https://huggingface.co/black-forest-labs/FLUX.1-schnell
+pub const FLUX_SCHNELL_AUTOENCODER : WeightDescriptor = weight!(
+  "Flux-1 Schnell (Autoencoder)",
+  "flux1-schnell.ae.safetensors",
+  WeightFunction::ImageGeneration
+);
+
+/// https://huggingface.co/lmz/candle-flux
+pub const LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF : WeightDescriptor = weight!(
+  "LMZ Flux Schnell Quantized",
+  "lmz_candle_flux_flux1-schnell.gguf",
+  WeightFunction::ImageGeneration
+);
+
+/// https://huggingface.co/google/t5-v1_1-xxl/tree/main
+/// NB: I couldn't find the safetensors file here
+/// NB: This should live under "refs/pr/2"
+pub const GOOGLE_T5_V1_1_XXL_MODEL : WeightDescriptor = weight!(
+  "Google T5 v1.1 (XXL)",
+  "google_t5_v1_1_xxl_model.safetensors",
+  WeightFunction::TextEncoder
+);
+
+/// https://huggingface.co/google/t5-v1_1-xxl/tree/main
+/// NB: I couldn't find the safetensors file here
+/// NB: This should live under "refs/pr/2"
+pub const GOOGLE_T5_V1_1_XXL_CONFIG : WeightDescriptor = weight!(
+  "Google T5 v1.1 (XXL) Config",
+  "google_t5_v1_1_xxl_config.json",
+  WeightFunction::TextEncoder
+);
+
+/// https://huggingface.co/lmz/mt5-tokenizers
+pub const LMZ_T5_TOKENIZER_JSON : WeightDescriptor = weight!(
+  "LMZ T5 Tokenizer",
+  "lmz_mt5_tokenizers_t5-v1_1-xxl.tokenizer.json",
+  WeightFunction::TextTokenizer
+);
+
+/// https://huggingface.co/openai/clip-vit-large-patch14/tree/main
+pub const OPENAI_CLIP_VIT_P14_MODEL : WeightDescriptor = weight!(
+  "OpenAI CLIP VIT (patch 14) model",
+  "openai_clip-vit-large-patch14.model.safetensors",
+  WeightFunction::TextTokenizer
+);
+
+/// https://huggingface.co/openai/clip-vit-large-patch14/tree/main
+pub const OPENAI_CLIP_VIT_P14_TOKENIZER_JSON: WeightDescriptor = weight!(
+  "OpenAI CLIP VIT (patch 14) tokenizer",
+  "openai_clip-vit-large-patch14.tokenizer.json",
+  WeightFunction::TextTokenizer
+);
@@ -0,0 +1,4 @@
+mod weight_descriptor_builder;
+pub mod flux_weights;
+pub mod weight_descriptor;
+pub mod weights;
@@ -0,0 +1,47 @@
+
+// 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 {
+  ImageGeneration,
+  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::ImageGeneration => "image generation",
+      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
+);
@@ -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

@@ -20,29 +21,35 @@ tauri-build = { version = "2.0.4", features = [] }
 [dependencies]

 # Internal
+crockford = { path = "../../lib/data/crockford" }
 filesys = { path = "../../lib/files/filesys" }
 memory_caching = { path = "../../lib/caching/memory_caching" }
+ml_weights_registry = { path = "../ml_weights_registry" }
+primitives = { path = "../../lib/data/primitives" }

+# Internal ML
+# TODO: Only on Windows and Linux
+#ml_models = { path = "../ml_models", optional = true }
 # Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332
-cudarc = "0.13.7"
+cudarc = { version = "0.13.5", features = ["std", "cublas", "cublaslt", "curand", "driver", "nvrtc", "f16", "cuda-version-from-build-system", "dynamic-linking"], default-features=false }

 # External
 #openssl = { version = "0.10", features = ["vendored"] }
+# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls
 anyhow = "1.0"
 base64 = "0.22.1"
 bytes = "1.10"
-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"] }
-dirs = "6.0.0"
+chrono = "0.4.40"
+const_format = "0.2.34"
 directories = "6.0.0"
+dirs = "6.0.0"
 futures = "0.3"
-hf-hub = "0.3.2"
+half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
 image = "0.25.5"
 log = "0.4"
 once_cell = "1.20.3"
+ort = { version = "1.16.2", features = ["cuda", "load-dynamic"] }
 rand = "0.8.5"
-# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls
 reqwest = { version = "0.12.12", features = ["stream"], default-features = false }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
@@ -51,25 +58,22 @@ tauri = { version = "2.2.4", features = [] }
 tauri-plugin-log = "2"
 tempdir = "0.3.7"
 tempfile = "3.18.0"
-tokenizers = "0.15.2"
 tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] }
 tokio-util = { version = "0.7", features = ["codec"] }
-half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }

-[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"
-]
+# 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.
@@ -1,16 +1,12 @@
-use crate::ml::downloads::download_all_models::download_all_models;
 use crate::state::app_dir::AppDataRoot;
 use tauri::State;

 #[tauri::command]
 pub async fn download_models(
-  app_data_root: State<'_, AppDataRoot>,
+  _app_data_root: State<'_, AppDataRoot>,
 ) -> Result<String, String> {
-  
-  download_all_models(&app_data_root).await
-    .map_err(|err| {
-      "error downloading models".to_string()
-    })?;
-  
+
+  // NB: This endpoint has been removed for now as we now download models in the app lifecycle.
+
  Ok(("downloaded".to_string()))
 }
@@ -1,146 +0,0 @@
-use crate::ml::downloads::download_all_models::download_all_models;
-use crate::ml::model_cache::ModelCache;
-use crate::ml::prompt_cache::PromptCache;
-use crate::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
-use crate::state::app_config::AppConfig;
-use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
-use bytes::BytesMut;
-use image::imageops::FilterType;
-use image::{DynamicImage, ImageFormat, ImageReader};
-use log::{error, info};
-use std::io::Cursor;
-use std::path::PathBuf;
-use tauri::{AppHandle, State};
-use crate::state::app_dir::AppDataRoot;
-
-const PROMPT_FILENAME : &str = "prompt.txt";
-const PROMPT: &str = "green goblin, blood dripping from lips, detailed, photorealistic, 8k";
-
-const NEGATIVE_PROMPT: &str = "bad quality, bad faces, poor quality, blurry faces, watermark";
-
-const RANDOM_SEED: u32 = 42;
-
-// TODO: Mark async https://github.com/tauri-apps/tauri/discussions/7737
-
-
-/// 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<u8>,
-  model_config: State<'_, AppConfig>,
-  model_cache: State<'_, ModelCache>,
-  prompt_cache: State<'_, PromptCache>,
-  app_data_root: State<'_, AppDataRoot>,
-  app: AppHandle,
-) -> Result<String, String> {
-
-  let prompt = get_prompt_or_fallback(prompt);
-
-  info!("Prompt: {}", prompt);
-  info!("Strength: {:?}", strength);
-
-  let image = hydrate_base64_image(image)
-    .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
-
-  // TODO(bt): Move this to another endpoint
-  download_all_models(&app_data_root).await
-    .map_err(|err| {
-      error!("couldn't download models: {:?}", err);
-      "Couldn't download models".to_string()
-    })?;
-
-  let result = do_infer_image(
-    &prompt, 
-    image, 
-    strength, 
-    &model_config, 
-    &model_cache, 
-    prompt_cache, 
-    app, 
-    &app_data_root
-  ).await;
-  
-  if let Err(err) = result.as_deref() {
-    error!("There was an error: {:?}", err);
-  }
-  
-  result
-}
-
-async fn do_infer_image(
-  prompt: &str,
-  image: DynamicImage,
-  strength: Option<u8>,
-  config: &AppConfig,
-  model_cache: &ModelCache,
-  prompt_cache: State<'_, PromptCache>,
-  app: AppHandle,
-  app_data_root: &AppDataRoot,
-) -> Result<String, String> {
-
-  let args = Args {
-    image: &image,
-    prompt: prompt.to_string(),
-    uncond_prompt: NEGATIVE_PROMPT.to_string(),
-    model_cache,
-    configs: config,
-    prompt_cache: &prompt_cache,
-    i2i_strength: strength,
-    cfg_scale: config.cfg_scale,
-    app: &app,
-    app_data_root,
-    use_flash_attn: true,
-  };
-
-  match lcm_pipeline(args) {
-    Ok(image) => {
-      let mut bytes = Vec::with_capacity(1024*1024);
-      
-      image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)
-        .map_err(|err| format!("failure to encode image: {:?}", err))?;
-
-      let bytes = BASE64_STANDARD.encode(bytes);
-
-      Ok(bytes)
-    },
-    Err(e) => Err(format!("Failed to generate image: {}", e)),
-  }
-}
-
-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()
-}
-
-fn hydrate_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()?;
-
-  // TODO(bt,2025-02-17): Running out of vram with full image buffer size
-  let image = image.resize(512, 512, FilterType::CatmullRom);
-
-  Ok(image)
-}
@@ -0,0 +1,13 @@
+use log::info;
+
+
+#[tauri::command]
+pub async fn inpaint_image(
+  image: &str,
+  mask: &str,
+  prompt: String,
+) -> Result<String, String> {
+  info!("inpaint_endpoint endpoint called.");
+
+  Ok(mask.to_string())
+}
@@ -1,4 +1,7 @@
 pub mod download_models;
 pub mod flip_image;
-pub mod image_endpoint;
-pub mod test_counter;
+pub mod inpaint_image_endpoint;
+pub mod realtime_image_endpoint;
+pub mod remove_background_endpoint;
+pub mod save_image_endpoint;
+pub mod text_to_image_endpoint;
@@ -0,0 +1,148 @@
+#[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>,
+  seed: Option<u64>,
+  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);
+
+  let seed = seed.or(model_config.seed);
+
+  #[cfg(not(target_os = "macos"))]
+  {
+    let result = infer_image_impl(
+      &prompt,
+      image,
+      strength,
+      seed,
+      &model_config,
+      &model_cache,
+      prompt_cache,
+      &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>,
+  seed: Option<u64>,
+  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: seed,
+    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,38 @@
+use crate::state::app_config::AppConfig;
+use crate::state::app_dir::AppDataRoot;
+use crate::stubs::model_cache::ModelCache;
+use crate::stubs::prompt_cache::PromptCache;
+use crate::utils::image::decode_base64_image::decode_base64_image;
+use chrono::Local;
+use crockford::crockford_entropy_lower;
+use image::{ImageFormat, RgbImage};
+use log::info;
+use tauri::{AppHandle, State};
+
+const FILENAME_LENGTH : usize = 10;
+
+/// This handler takes an image (as a base64 encoded string) and saves it.
+#[tauri::command]
+pub async fn save_image(
+  image: &str,
+  app_data_root: State<'_, AppDataRoot>,
+) -> Result<String, String> {
+  info!("save_image endpoint called.");
+
+  let image = decode_base64_image(image)
+    .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
+
+  let assets_dir = app_data_root.assets_dir();
+  
+  let date_save_dir = assets_dir.make_or_get_current_date_dir()
+    .map_err(|err| format!("Couldn't create assets save directory: {}", err))?;
+
+  let entropy = crockford_entropy_lower(FILENAME_LENGTH); // TODO: These are ugly. A monotonic counter would be better.
+  let filename = format!("{}.png", entropy);
+  let filename = date_save_dir.join(filename);
+  
+  image.save_with_format(filename,  ImageFormat::Png)
+    .map_err(|err| format!("Couldn't save image: {}", err))?;
+
+  Ok("".to_string())
+}
@@ -1,24 +0,0 @@
-use std::sync::{Arc, RwLock};
-use once_cell::sync::Lazy;
-
-static COUNTER : Lazy<Arc<RwLock<u64>>> = Lazy::new(|| Arc::new(RwLock::new(0)));
-
-#[tauri::command]
-pub fn test_counter() -> u64 {
-  let value : u64;
-  {
-    match COUNTER.write() {
-      Ok(mut counter) => {
-        *counter += 1;
-        value = *counter;
-      },
-      Err(_e) => {
-        value = 0;
-      }
-    }
-  }
-
-  println!("I was invoked from JavaScript! {:?}", value);
-
-  value
-}
@@ -0,0 +1,114 @@
+#[cfg(not(target_os = "macos"))]
+use ml_models::ml::flux::flux_pipeline::{flux_pipeline, FluxModel, FluxPipelineArgs};
+
+use crate::state::app_config::AppConfig;
+use crate::state::app_dir::{AppDataRoot, AppWeightsDir};
+use crate::stubs::model_cache::ModelCache;
+use crate::stubs::prompt_cache::PromptCache;
+use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png;
+use image::imageops::FilterType;
+use image::{DynamicImage, ImageReader, RgbImage};
+use log::{error, info};
+
+use ml_weights_registry::weights_registry::flux_weights::{FLUX_DEV, FLUX_SCHNELL, FLUX_SCHNELL_AUTOENCODER, GOOGLE_T5_V1_1_XXL_CONFIG, GOOGLE_T5_V1_1_XXL_MODEL, LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF, LMZ_T5_TOKENIZER_JSON, OPENAI_CLIP_VIT_P14_MODEL, OPENAI_CLIP_VIT_P14_TOKENIZER_JSON};
+use once_cell::sync::Lazy;
+use std::io::Cursor;
+use tauri::{AppHandle, State};
+
+const 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: AppHandle,
+  app_data_root: State<'_, AppDataRoot>,
+) -> 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_data_root.weights_dir(),
+      app,
+    ).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>,
+  weights_dir: &AppWeightsDir,
+  app: AppHandle,
+) -> Result<RgbImage, String> {
+  
+  let args = FluxPipelineArgs {
+    image: &image,
+    prompt: prompt.to_string(),
+    model_cache,
+    prompt_cache: &prompt_cache,
+    app: &app,
+    flux_model: FluxModel::Schnell,
+    active_blocks: None,
+    use_cuda_stream: Some(true),
+    prefetch_next_batch: Some(false),
+    use_full_gpu: Some(true),
+    img2img_strength: strength,
+    seed: None,
+    use_quantized_model: Some(true),
+    model_config: &config.model_config,
+    flux_dev_or_schnell_path: &weights_dir.weight_path(&FLUX_SCHNELL),
+    flux_schnell_quantized_gguf: &weights_dir.weight_path(&LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF),
+    flux_autoencoder_path: &weights_dir.weight_path(&FLUX_SCHNELL_AUTOENCODER),
+    google_t5_model_path: &weights_dir.weight_path(&GOOGLE_T5_V1_1_XXL_MODEL),
+    google_t5_config_path: &weights_dir.weight_path(&GOOGLE_T5_V1_1_XXL_CONFIG),
+    lmz_t5_tokenizer_path: &weights_dir.weight_path(&LMZ_T5_TOKENIZER_JSON),
+    openai_clip_model_path: &weights_dir.weight_path(&OPENAI_CLIP_VIT_P14_MODEL),
+    openai_clip_tokenizer_json_path: &weights_dir.weight_path(&OPENAI_CLIP_VIT_P14_TOKENIZER_JSON),
+  };
+
+  let image = flux_pipeline(args)
+    .map_err(|err| format!("failure to encode image: {:?}", err))?;
+
+  Ok(image)
+}
@@ -1,13 +1,16 @@
 use serde::Serialize;

-
 #[derive(Clone, Copy, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub enum NotificationModelType {
+  ImageSegmentation,
+  TextEncoder,
+  TextTokenizer,
  Unet,
  Vae,
  Json,
  ClipEncoder,
+  FluxModel,
 }

 #[derive(Clone, Serialize)]
@@ -26,5 +29,17 @@ pub enum NotificationEvent<'a> {
  ModelDownloadComplete {
    model_name: &'a str,
    model_type: NotificationModelType,
-  }
+  },
+  ModelLoading {
+    model_name: String,
+    model_type: NotificationModelType,
+  },
+  ModelInference {
+    model_name: String,
+    model_type: NotificationModelType,
+  },
+  ModelInferenceComplete {
+    model_name: String,
+    model_type: NotificationModelType,
+  },
 }
@@ -1,17 +1,23 @@
 pub mod endpoints;
 pub mod events;
-pub mod ml;
 pub mod state;
+pub mod stubs;
 pub mod threads;
 pub mod transfer;
+pub mod utils;

 use crate::endpoints::download_models::download_models;
-use crate::endpoints::image_endpoint::infer_image;
-use crate::endpoints::test_counter::test_counter;
-use crate::ml::model_cache::ModelCache;
-use crate::ml::prompt_cache::PromptCache;
+use crate::endpoints::flip_image::flip_image;
+use crate::endpoints::inpaint_image_endpoint::inpaint_image;
+use crate::endpoints::realtime_image_endpoint::infer_image;
+use crate::endpoints::remove_background_endpoint::remove_background;
+use crate::endpoints::save_image_endpoint::save_image;
+use crate::endpoints::text_to_image_endpoint::text_to_image;
 use crate::state::app_config::AppConfig;
+use crate::stubs::model_cache::ModelCache;
+use crate::stubs::prompt_cache::PromptCache;
 use crate::threads::downloader_thread::downloader_thread;
+
 use tauri_plugin_log::Target;
 use tauri_plugin_log::TargetKind;

@@ -39,9 +45,15 @@ pub fn run() {
  tauri::Builder::default()
    .plugin(tauri_plugin_log::Builder::new()
      .level(log::LevelFilter::Info)
-      .targets(vec![Target::new(TargetKind::Stdout)])
+      .targets(vec![
+        Target::new(TargetKind::Stdout),
+        Target::new(TargetKind::LogDir { file_name: Some(app_data_root.log_file_name_str().to_string()) }),
+      ])
      .build())
    .setup(|app| {
+      // TODO(bt): This is broken on windows
+      // log_environment_details();
+
      //if cfg!(debug_assertions) {
      //  app.handle().plugin(
      //    tauri_plugin_log::Builder::default()
@@ -60,9 +72,14 @@ pub fn run() {
    .manage(model_cache)
    .manage(app_data_root)
    .invoke_handler(tauri::generate_handler![
-      test_counter, 
-      infer_image,
      download_models,
+      flip_image,
+      infer_image,
+      inpaint_image,
+      remove_background,
+      save_image,
+      text_to_image,
+      text_to_image,
    ])
    .run(tauri::generate_context!())
    .expect("error while running tauri application");
@@ -1,15 +0,0 @@
-use crate::ml::downloads::download_model_registry_file::download_model_registry_file;
-use crate::ml::model_type::ModelType;
-use crate::state::app_dir::AppDataRoot;
-
-pub async fn download_all_models(app_data_root: &AppDataRoot) -> anyhow::Result<()> {
-  let weights_dir = app_data_root.weights_dir();
-  
-  //download_model_registry_file(ModelType::ClipJson, weights_dir).await?;
-  //download_model_registry_file(ModelType::SdxlTurboUnet, weights_dir).await?;
-  //download_model_registry_file(ModelType::SdxlTurboVae, weights_dir).await?;
-  //download_model_registry_file(ModelType::SdxlTurboClipEncoder, weights_dir).await?;
-  //download_model_registry_file(ModelType::SdxlTurboClipEncoder2, weights_dir).await?;
-  
-  Ok(())
-}
@@ -1,30 +0,0 @@
-use crate::ml::model_type::ModelType;
-use crate::state::app_dir::AppWeightsDir;
-use log::info;
-use std::fs::File;
-use std::path::{Path, PathBuf};
-
-pub async fn download_model_registry_file(
-  model_type: ModelType,
-  weights_dir: &AppWeightsDir,
-) -> anyhow::Result<PathBuf> {
-  
-  let path = weights_dir.model_path(&model_type);
-
-  if path.exists() {
-    return Ok(path);
-  }
-
-  info!("downloading model: {:?} to {:?}", model_type, path);
-
-  let body = reqwest::get(model_type.get_download_url())
-    .await?
-    .bytes()
-    .await?;
-  
-  let _file = File::create(&path)?;
-  
-  std::fs::write(&path, body.as_ref())?;
-  
-  Ok(path)
-}
@@ -1,2 +0,0 @@
-pub mod download_model_registry_file;
-pub mod download_all_models;
@@ -1,8 +0,0 @@
-pub mod image;
-pub mod model_cache;
-pub mod model_file;
-pub mod models;
-pub mod prompt_cache;
-pub mod stable_diffusion;
-pub mod model_type;
-pub mod downloads;
@@ -1,86 +0,0 @@
-use crate::events::notification_event::NotificationModelType;
-use crate::state::app_dir::AppWeightsDir;
-use std::path::PathBuf;
-
-#[derive(Copy, Clone, Debug)]
-pub enum ModelType {
-  ClipJson,
-  SdxlTurboUnet,
-  SdxlTurboVae,
-  SdxlTurboClipEncoder,
-  SdxlTurboClipEncoder2,
-  SimianLuoLcmDreamshaperV7Unet,
-  LykonDreamshaper7Vae,
-  LykonDreamshaper7TextEncoderFp16,
-}
-
-impl ModelType {
-  pub fn get_name(&self) -> &'static str {
-    match self {
-      Self::ClipJson => "CLIP (JSON)",
-      Self::SdxlTurboUnet => "SDXL Turbo UNET",
-      Self::SdxlTurboVae => "SDXL Turbo VAE",
-      Self::SdxlTurboClipEncoder => "Clip Encoder #1",
-      Self::SdxlTurboClipEncoder2 => "Clip Encoder #2",
-      Self::SimianLuoLcmDreamshaperV7Unet => "SimianLuo LCM Dreamshaper v7 UNET",
-      Self::LykonDreamshaper7Vae => "Lykon Dreamshaper 7 VAE",
-      Self::LykonDreamshaper7TextEncoderFp16 => "Lykon Dreamshaper Text Encoder (fp16)",
-    }
-  }
-
-  pub fn get_notification_type(&self) -> NotificationModelType {
-    match self {
-      Self::ClipJson => NotificationModelType::Json,
-      Self::SdxlTurboUnet => NotificationModelType::Unet,
-      Self::SdxlTurboVae => NotificationModelType::Vae,
-      Self::SdxlTurboClipEncoder => NotificationModelType::ClipEncoder,
-      Self::SdxlTurboClipEncoder2 => NotificationModelType::ClipEncoder,
-      Self::SimianLuoLcmDreamshaperV7Unet => NotificationModelType::Unet,
-      Self::LykonDreamshaper7Vae => NotificationModelType::Vae,
-      Self::LykonDreamshaper7TextEncoderFp16 => NotificationModelType::ClipEncoder,
-    }
-  }
-
-  pub fn get_hf_id(&self) -> Option<&'static str> {
-    match self {
-      Self::ClipJson => Some("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"),
-      Self::SdxlTurboUnet 
-      | Self::SdxlTurboVae
-      | Self::SdxlTurboClipEncoder
-      | Self::SdxlTurboClipEncoder2 => Some("stabilityai/sdxl-turbo"),
-      Self::SimianLuoLcmDreamshaperV7Unet => Some("SimianLuo/LCM_Dreamshaper_v7"),
-      Self::LykonDreamshaper7Vae
-      | Self::LykonDreamshaper7TextEncoderFp16 => Some("Lykon/dreamshaper-7"),
-    }
-  }
-
-  pub fn get_download_url(&self) -> &'static str {
-    match self {
-      Self::ClipJson => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/tokenizer.json",
-      Self::SdxlTurboUnet => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/diffusion_pytorch_model.unet.safetensors",
-      Self::SdxlTurboVae => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/diffusion_pytorch_model.vae.safetensors",
-      Self::SdxlTurboClipEncoder => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/clip_text_encoder.safetensors",
-      Self::SdxlTurboClipEncoder2 => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/clip_text_encoder_2.safetensors",
-      Self::SimianLuoLcmDreamshaperV7Unet => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/simianluo_lcm_dreamshaper_v7_unet.safetensors",
-      Self::LykonDreamshaper7Vae => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_vae.safetensors",
-      Self::LykonDreamshaper7TextEncoderFp16 => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_text_encoder.fp16.safetensors",
-    }
-  }
-
-  pub fn get_filename(&self) -> &'static str {
-    match self {
-      Self::ClipJson => "clip_vit_bigg_14_lion2b_39b_b160k.tokenizer.json",
-      Self::SdxlTurboUnet => "diffusion_pytorch_model.unet.safetensors",
-      Self::SdxlTurboVae => "diffusion_pytorch_model.vae.safetensors",
-      Self::SdxlTurboClipEncoder => "clip_text_encoder.safetensors",
-      Self::SdxlTurboClipEncoder2 => "clip_text_encoder_2.safetensors",
-      Self::SimianLuoLcmDreamshaperV7Unet => "simianluo_lcm_dreamshaper_v7_unet.safetensors",
-      Self::LykonDreamshaper7Vae => "lykon_dreamshaper_7_vae.safetensors",
-      Self::LykonDreamshaper7TextEncoderFp16 => "lykon_dreamshaper_7_text_encoder.fp16.safetensors",
-    }
-  }
-  
-  pub fn get_path(&self, weights_dir: &AppWeightsDir) -> PathBuf {
-    weights_dir.model_path(self)
-  }
-}
@@ -1,130 +0,0 @@
-use crate::ml::model_file::StableDiffusionVersion;
-use anyhow::anyhow;
-use candle_core::{DType, Device, Tensor};
-use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
-use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
-use hf_hub::api::sync::Api;
-use std::sync::{Arc, RwLock};
-
-pub struct LazyLoadVaeModel {
-  lazy_model: Arc<RwLock<Option<AutoEncoderKL>>>,
-  model_configs: VaeConfigs,
-}
-
-// TODO: Maybe we pass these via a central config read at startup, passed as an Arc<ConfigT>
-pub struct VaeConfigs {
-  sd_config: StableDiffusionConfig,
-  sd_version: StableDiffusionVersion,
-  device: Device,
-  dtype: DType,
-  hf_api: Api,
-}
-
-impl LazyLoadVaeModel {
-  pub fn new(
-    sd_config: &StableDiffusionConfig,
-    sd_version: &StableDiffusionVersion,
-    hf_api: &Api,
-    device: &Device,
-    dtype: DType,
-  ) -> anyhow::Result<Self> {
-    println!("building vae model... (1)");
-    
-    Ok(Self {
-      lazy_model: Arc::new(RwLock::new(None)),
-      model_configs: VaeConfigs {
-        sd_config: sd_config.clone(),
-        sd_version: sd_version.clone(),
-        device: device.clone(),
-        hf_api: hf_api.clone(),
-        dtype
-      }
-    })
-  }
-
-  pub fn encode(&self, xs: &Tensor) -> anyhow::Result<DiagonalGaussianDistribution> {
-    println!("vae encode...");
-    match self.try_encode(xs) {
-      Err(err) => Err(err),
-      Ok(Some(result)) => Ok(result),
-      Ok(None) => {
-        // Model wasn't previously loaded
-        self.load_model()?;
-        self.try_encode(xs)?
-          .ok_or(anyhow!("model was not previously loaded"))
-      }
-    }
-  }
-  
-  pub fn decode(&self, xs: &Tensor) -> anyhow::Result<Tensor> {
-    println!("vae decode...");
-    match self.try_decode(xs) {
-      Err(err) => Err(err),
-      Ok(Some(result)) => Ok(result),
-      Ok(None) => {
-        // Model wasn't previously loaded
-        self.load_model()?;
-        self.try_decode(xs)?
-          .ok_or(anyhow!("model was not previously loaded"))
-      }
-    }
-  }
-
-  fn try_encode(&self, xs: &Tensor) -> anyhow::Result<Option<DiagonalGaussianDistribution>> {
-    match self.lazy_model.read() {
-      Err(err) => Err(anyhow!("lock error: {:?}", err)),
-      Ok(model) => {
-        match &*model {
-          None => Ok(None), // Model is not yet loaded
-          Some(model) => {
-            Ok(Some(model.encode(xs)?))
-          }
-        }
-      }
-    }
-  }
-
-  fn try_decode(&self, xs: &Tensor) -> anyhow::Result<Option<Tensor>> {
-    match self.lazy_model.read() {
-      Err(err) => Err(anyhow!("lock error: {:?}", err)),
-      Ok(model) => {
-        match &*model {
-          None => Ok(None), // Model is not yet loaded
-          Some(model) => {
-            Ok(Some(model.decode(xs)?))
-          }
-        }
-      }
-    }
-  }
-
-  fn load_model(&self) -> anyhow::Result<()> {
-    println!("load vae");
-    match self.lazy_model.write() {
-      Err(err) => return Err(anyhow!("lock error: {:?}", err)),
-      Ok(mut maybe_model) => {
-        match &*maybe_model {
-          Some(_model) => {} // Already loaded; fall through
-          None => {
-            let repo = self.model_configs.sd_version.repo();
-
-            println!("Building VAE model from : {:?} ... (3)", repo);
-            
-            let vae_file = self.model_configs.hf_api.model(repo.to_string())
-              .get("vae/diffusion_pytorch_model.safetensors")?;
-
-            println!("Building VAE model from file {:?}...", &vae_file);
-            let vae = self.model_configs
-              .sd_config
-              .build_vae(vae_file, &self.model_configs.device, self.model_configs.dtype)?;
-
-            println!("Built VAE...");
-            *maybe_model = Some(vae);
-          }
-        }
-      }
-    }
-    Ok(())
-  }
- 
-}
@@ -1,2 +0,0 @@
-pub mod unet_model;
-pub mod lazy_load_vae_model;
@@ -1,30 +0,0 @@
-use candle_core::{DType, Device, Tensor};
-use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
-use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
-use std::path::Path;
-
-pub struct UNetModel {
-  model: UNet2DConditionModel,
-}
-
-impl UNetModel {
-  pub fn new<P: AsRef<Path>>(sd_config: &StableDiffusionConfig, unet_weights_file: P, device: &Device, dtype: DType) -> anyhow::Result<Self> {
-    println!("building unet model... (1)");
-
-    let model = sd_config.build_unet(unet_weights_file, &device, 4, true, dtype)?;
-
-    Ok(Self { model })
-  }
-
-  pub fn new_with_inner(model: UNet2DConditionModel) -> Self {
-    Self { model }
-  }
-
-  pub fn inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result<Tensor> {
-    Ok(self.model.forward(&latent_model_input, timestep, &text_embeddings)?)
-  }
-
-  pub fn forward_with_guidance(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor, guidance_scale_embedding: &Tensor) -> anyhow::Result<Tensor> {
-    Ok(self.model.forward_with_guidance(&latent_model_input, timestep, &text_embeddings, Some(guidance_scale_embedding))?)
-  }
-}
@@ -1,32 +1,14 @@
-use crate::ml::model_file::StableDiffusionVersion;
 use crate::state::app_dir::AppDataRoot;
 use crate::state::os_platform::OsPlatform;
 use crate::state::yaml_config::YamlConfig;
-use candle_core::{DType, Device};
-use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
-use filesys::recursively_find_file_by_name::recursively_find_file_by_name;
-use hf_hub::api::sync::Api;
-use log::info;
-use std::path::PathBuf;
+use crate::stubs::model_config::ModelConfig;

 const DEFAULT_SD_IMAGE_WIDTH: usize = 1024;
 const DEFAULT_SD_IMAGE_HEIGHT: usize = 1024;

-const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32;
-const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::V1_5;

 /// A central place to configure app-wide details.
 pub struct AppConfig {
-  /// We only support one device type currently, and we can't multiplex 
-  /// over several devices yet.
-  pub device: Device,
-  
-  /// The core numeric type to use for models.
-  pub dtype: DType,
-  
-  pub sd_version: StableDiffusionVersion,
-  pub sd_config: StableDiffusionConfig,
-
  pub image_height: usize,
  pub image_width: usize,
  pub scheduler_steps: usize,
@@ -36,11 +18,10 @@ pub struct AppConfig {

  pub cfg_scale: Option<f64>,
  
-  /// Probably shouldn't be here.
-  pub hf_api: Api,
-  
  /// Root location for application data.
  pub app_data_root: AppDataRoot,
+  
+  pub model_config: ModelConfig,
 }

 impl AppConfig {
@@ -50,30 +31,6 @@ impl AppConfig {
    
    println!("Configs: {:?}", yaml_configs);
    
-    let device = Device::new_cuda(0)
-      .unwrap_or_else(|e| {
-          info!("CUDA not available ({}), falling back to CPU", e);
-          Device::Cpu
-        });
-
-    
-    let sd_config = match SD_VERSION {
-      StableDiffusionVersion::Turbo => {
-        StableDiffusionConfig::sdxl_turbo(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
-      }
-      StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => {
-        StableDiffusionConfig::v1_5(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
-      }
-      StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint => {
-        StableDiffusionConfig::v2_1(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
-      }
-      StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint => {
-        StableDiffusionConfig::sdxl(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
-      }
-    };
-    
-    let hf_api = Api::new()?;
-    
    let os_platform = OsPlatform::get();

    let default_data_root = match os_platform {
@@ -92,18 +49,14 @@ impl AppConfig {
    println!("App data root: {:?}", app_data_root.path());

    Ok(Self {
-      device,
-      dtype: DEFAULT_SD_NUMERIC_DATATYPE,
-      sd_version: SD_VERSION,
-      sd_config,
      image_height: yaml_configs.image_height.unwrap_or(DEFAULT_SD_IMAGE_HEIGHT),
      image_width: yaml_configs.image_width.unwrap_or(DEFAULT_SD_IMAGE_WIDTH),
      scheduler_steps: yaml_configs.scheduler_steps.unwrap_or(4),
      scheduler_samples: yaml_configs.scheduler_samples.unwrap_or(15),
      seed: yaml_configs.seed,
      cfg_scale: yaml_configs.cfg_scale,
-      hf_api,
      app_data_root,
+      model_config: ModelConfig::init()?,
    })
  }
 }
@@ -1,17 +1,22 @@
-use crate::ml::model_type::ModelType;
 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 chrono::{DateTime, Local};
 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";

+/// Note: Tauri appends ".log" to the end of the filename.
+const LOG_FILE_NAME : &str = "application_debug";
+
 /// The path to the application data directory, which includes "asset" and "weights" data.
 #[derive(Clone)]
 pub struct AppDataRoot {
@@ -19,6 +24,8 @@ pub struct AppDataRoot {
  assets_dir: AppAssetsDir,
  weights_dir: AppWeightsDir,
  temp_dir: TemporaryDir,
+  log_file_name: PathBuf,
+  log_file_name_string: String,
 }

 #[derive(Clone)]
@@ -69,12 +76,19 @@ impl AppDataRoot {
    let assets_dir = AppAssetsDir::create_existing(dir.join(ASSETS_SUBDIRECTORY))?;
    let weights_dir = AppWeightsDir::create_existing(dir.join(WEIGHTS_SUBDIRECTORY))?;
    let temp_dir = TemporaryDir::create_existing(dir.join(TEMPORARY_SUBDIRECTORY))?;
+    let log_file_name = dir.join(LOG_FILE_NAME);
+    let log_file_name_string = log_file_name
+        .to_str()
+        .ok_or(anyhow!("couldn't convert log path to str"))?
+        .to_string();

    Ok(Self {
      path: dir,
      assets_dir,
      weights_dir,
      temp_dir,
+      log_file_name,
+      log_file_name_string,
    })
  }
  
@@ -93,6 +107,14 @@ impl AppDataRoot {
  pub fn path(&self) -> &Path {
    &self.path
  }
+
+  pub fn log_file_name(&self) -> &Path {
+    &self.log_file_name
+  }
+
+  pub fn log_file_name_str(&self) -> &str {
+    &self.log_file_name_string
+  }
 }

 impl AppAssetsDir {
@@ -116,6 +138,18 @@ impl AppAssetsDir {
  pub fn path(&self) -> &Path {
    &self.path
  }
+  
+  /// Return the current date asset directory
+  /// If it doesn't exist, create it.
+  pub fn make_or_get_current_date_dir(&self) -> anyhow::Result<PathBuf> {
+    let format: DateTime<Local> = Local::now();
+    let date_directory_name = format.format("%Y-%d-%m").to_string();
+    let full_path = self.path.join(date_directory_name);
+    if !full_path.exists() {
+      std::fs::create_dir(&full_path)?;
+    }
+    Ok(full_path)
+  }
 }

 impl AppWeightsDir {
@@ -139,9 +173,9 @@ impl AppWeightsDir {
  pub fn path(&self) -> &Path {
    &self.path
  }
-  
-  pub fn model_path(&self, model_registry: &ModelType) -> PathBuf {
-    self.path.join(model_registry.get_filename())
+
+  pub fn weight_path(&self, descriptor: &WeightDescriptor) -> PathBuf {
+    self.path.join(descriptor.filename)
  }
 }

@@ -0,0 +1,3 @@
+pub mod model_cache;
+pub mod model_config;
+pub mod prompt_cache;
@@ -0,0 +1,14 @@
+#[cfg(not(target_os = "macos"))]
+pub use ml_models::ml::model_cache::ModelCache;
+
+#[cfg(not(not(target_os = "macos")))]
+pub struct ModelCache {
+  // Intentionally left blank
+}
+
+#[cfg(not(not(target_os = "macos")))]
+impl ModelCache {
+  pub fn new() -> Self {
+    Self {}
+  }
+}
@@ -0,0 +1,14 @@
+#[cfg(not(target_os = "macos"))]
+pub use ml_models::ml::model_config::ModelConfig;
+
+#[cfg(not(not(target_os = "macos")))]
+pub struct ModelConfig {
+  // Intentionally left blank
+}
+
+#[cfg(not(not(target_os = "macos")))]
+impl ModelConfig {
+  pub fn init() -> anyhow::Result<Self> {
+    Ok(Self {})
+  }
+}
@@ -0,0 +1,15 @@
+
+#[cfg(not(target_os = "macos"))]
+pub use ml_models::ml::prompt_cache::PromptCache;
+
+#[cfg(not(not(target_os = "macos")))]
+pub struct PromptCache {
+  // Intentionally left blank
+}
+
+#[cfg(not(not(target_os = "macos")))]
+impl PromptCache {
+  pub fn with_capacity(capacity: usize) -> anyhow::Result<Self> {
+    Ok(Self {})
+  }
+}
@@ -1,10 +1,11 @@
 use crate::events::notification_event::{NotificationEvent, NotificationModelType};
-use crate::ml::model_type::ModelType;
 use crate::state::app_dir::AppDataRoot;
 use crate::transfer::download::{download_async, ProgressUpdate};
-use futures::stream::FuturesUnordered;
 use futures::StreamExt;
+use futures::stream::FuturesUnordered;
 use log::{debug, error, info};
+use ml_weights_registry::weights_registry::weight_descriptor::WeightFunction;
+use ml_weights_registry::weights_registry::weights::{CLIP_JSON, DIS_MEDIUM_ONNX, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2, SDXL_TURBO_UNET, SDXL_TURBO_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
 use std::collections::VecDeque;
 use std::path::PathBuf;
 use std::sync::mpsc::Receiver;
@@ -12,6 +13,7 @@ use std::sync::{Arc, Mutex};
 use tauri::{AppHandle, Emitter};
 use tempfile::NamedTempFile;
 use tokio::task::JoinHandle;
+use ml_weights_registry::weights_registry::flux_weights::{FLUX_SCHNELL, FLUX_SCHNELL_AUTOENCODER, GOOGLE_T5_V1_1_XXL_CONFIG, GOOGLE_T5_V1_1_XXL_MODEL, LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF, LMZ_T5_TOKENIZER_JSON, OPENAI_CLIP_VIT_P14_MODEL, OPENAI_CLIP_VIT_P14_TOKENIZER_JSON};

 const NOTIFICATION_CHANNEL_NAME : &str = "notification";
 const MAX_FILES : usize = 8;
@@ -39,22 +41,32 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> !

  let mut download_queue = VecDeque::new();

-  // TODO: Automatic enqueue of known important models + enqueue new models on-demand
-  download_queue.push_back(ModelType::ClipJson);
-  //download_queue.push_back(ModelType::SdxlTurboUnet);
-  //download_queue.push_back(ModelType::SdxlTurboVae);
-  download_queue.push_back(ModelType::SdxlTurboClipEncoder); // TODO(bt): Why is this still needed?
-  //download_queue.push_back(ModelType::SdxlTurboClipEncoder2);
-  download_queue.push_back(ModelType::SimianLuoLcmDreamshaperV7Unet);
-  download_queue.push_back(ModelType::LykonDreamshaper7Vae);
-  download_queue.push_back(ModelType::LykonDreamshaper7TextEncoderFp16);
+  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
+
+  // Flux:
+  download_queue.push_back(FLUX_SCHNELL);
+  download_queue.push_back(FLUX_SCHNELL_AUTOENCODER);
+  download_queue.push_back(LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF);
+  download_queue.push_back(GOOGLE_T5_V1_1_XXL_MODEL);
+  download_queue.push_back(GOOGLE_T5_V1_1_XXL_CONFIG);
+  download_queue.push_back(LMZ_T5_TOKENIZER_JSON);
+  download_queue.push_back(OPENAI_CLIP_VIT_P14_MODEL);
+  download_queue.push_back(OPENAI_CLIP_VIT_P14_TOKENIZER_JSON);

  let mut handles = FuturesUnordered::new();

  while let Some(model) = download_queue.pop_front() {
-    let url = model.get_download_url().to_string();
+    let url = model.r2_download_url.to_string();

-    let filename = app_data_root.weights_dir().model_path(&model);
+    let filename = app_data_root.weights_dir().weight_path(&model);

    if filename.exists() {
      info!("File already exists: {:?}", &filename);
@@ -75,9 +87,18 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> !
      }
    };

+    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,
+      WeightFunction::ImageGeneration => NotificationModelType::ImageSegmentation, // TODO FIX
+    };
+
    let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadStarted {
-      model_name: model.get_name(),
-      model_type: model.get_notification_type(),
+      model_name: model.name,
+      model_type,
    });

    if let Err(err) = result {
@@ -102,8 +123,8 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> !
        lock.push(ModelDownloadStatusChannel {
          receiver: rx,
          final_file_path: filename.clone(),
-          model_name: model.get_name().to_string(),
-          model_type: model.get_notification_type(),
+          model_name: model.to_descriptive_name(),
+          model_type,
        })
      }
    }
@@ -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;
@@ -0,0 +1,13 @@
+use actix_cors::Cors;
+use crate::util::netlify_branch_domain_matches::netlify_branch_domain_matches;
+
+pub fn add_artcraft_2d(cors: Cors, _is_production: bool) -> Cors {
+  cors
+      // Hypothetical domains
+      .allowed_origin("https://2d.storyteller.ai")
+      .allowed_origin("https://2d.getartcraft.com")
+      // Netlify project
+      .allowed_origin_fn(|origin, _req_head| {
+        netlify_branch_domain_matches(origin, "storyteller-2d.netlify.app")
+      })
+}
@@ -0,0 +1,13 @@
+use actix_cors::Cors;
+use crate::util::netlify_branch_domain_matches::netlify_branch_domain_matches;
+
+pub fn add_artcraft_3d(cors: Cors, _is_production: bool) -> Cors {
+  cors
+      // Hypothetical domains
+      .allowed_origin("https://3d.storyteller.ai")
+      .allowed_origin("https://3d.getartcraft.com")
+      // Netlify project
+      .allowed_origin_fn(|origin, _req_head| {
+        netlify_branch_domain_matches(origin, "storyteller-3d.netlify.app")
+      })
+}
@@ -1,3 +1,5 @@
+pub (crate) mod artcraft_2d;
+pub (crate) mod artcraft_3d;
 pub (crate) mod development_only;
 pub (crate) mod fakeyou;
 pub (crate) mod legacy;
@@ -2,7 +2,8 @@ use actix_cors::Cors;
 use log::info;

 use reusable_types::server_environment::ServerEnvironment;
-
+use crate::configs::artcraft_2d::add_artcraft_2d;
+use crate::configs::artcraft_3d::add_artcraft_3d;
 use crate::configs::development_only::add_development_only;
 use crate::configs::fakeyou::{add_fakeyou, add_fakeyou_dev_proxy};
 use crate::configs::legacy::{add_legacy_storyteller_stream, add_legacy_trumped, add_legacy_vocodes, add_power_stream};
@@ -40,6 +41,10 @@ fn do_build_cors_config(is_production: bool) -> Cors {
  cors = add_storyteller_board(cors, is_production);
  cors = add_storyteller_render(cors, is_production);

+  // Artcraft
+  cors = add_artcraft_2d(cors, is_production);
+  cors = add_artcraft_3d(cors, is_production);
+
  // Legacy
  cors = add_legacy_trumped(cors, is_production);
  cors = add_power_stream(cors, is_production);
@@ -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);
+  }
+}
@@ -0,0 +1,30 @@
+[package]
+name = "image_algorithms"
+edition = "2018"
+version = "0.0.1"
+authors = [
+    "Brandon Thomas <bt@brand.io>",
+    "Brandon Thomas <echelon@gmail.com>",
+]
+publish = false
+
+[lib]
+name = "image_algorithms"
+path = "src/lib.rs"
+
+[package.metadata.cargo-udeps.ignore]
+normal = ["my-workspace-hack"]
+
+[dependencies]
+
+# Workspace hack (faster compile times)
+my-workspace-hack = { version = "0.1", path = "../../my-workspace-hack" }
+
+# Internal
+errors = { path = "../errors" }
+
+# External
+image = "0.25.6"
+
+[dev-dependencies]
+testing = { path = "../testing" }
@@ -0,0 +1 @@
+pub mod remove_image_border;
@@ -0,0 +1,78 @@
+use image::{DynamicImage, Pixel};
+use errors::AnyhowResult;
+
+pub fn remove_image_border(image: &DynamicImage) -> AnyhowResult<DynamicImage> {
+  let mut image = image.to_rgba8();
+
+  let mut top = 0u32;
+  let mut bottom = image.height();
+  let mut left = 0u32;
+  let mut right = image.width();
+
+  // TODO: Iterate with binary search - it'll probably be much faster.
+
+  for (i, mut row) in image.rows().enumerate() {
+    let count_non_transparent = row.filter(|pixel| pixel.0[3] != 0).count();
+    if count_non_transparent < 5 {
+      top = i as u32;
+    } else {
+      break;
+    }
+  }
+
+  for (i, mut row) in image.rows().enumerate().rev() {
+    //let all_transparent = row.all(|pixel| pixel.0[3] == 0);
+    let count_non_transparent = row.filter(|pixel| pixel.0[3] != 0).count();
+    //if all_transparent {
+    if count_non_transparent < 5 {
+      bottom = i as u32;
+    } else {
+      break;
+    }
+  }
+
+  for (x, y, pixel) in image.enumerate_pixels() {
+    
+  }
+
+  for (y, mut row) in image.enumerate_rows() {
+    for (x, y, pixel) in row {
+      
+    }
+  }
+
+
+  let image = DynamicImage::ImageRgba8(image);
+
+  let width = image.width();
+  let height = bottom - top;
+
+  println!("top: {},  bottom: {}", top, bottom);
+  println!("orig width: {},  height: {}", image.width(), image.height());
+  println!("new width: {},  height: {}", width, height);
+
+  let image = image.crop_imm(0, top, width, height);
+
+  Ok(image)
+}
+
+#[cfg(test)]
+mod tests {
+  use image::{ImageFormat, ImageReader};
+  use errors::AnyhowResult;
+  use testing::test_file_path::test_file_path;
+  use crate::remove_image_border::remove_image_border;
+
+  #[ignore] // NB: Not going to check these files into the repo. Just for manual testing.
+  #[test]
+  fn test_crop_image() -> AnyhowResult<()> {
+    let filename = test_file_path("test_data/image/mario.png")?;
+    let image = ImageReader::open(filename)?.decode()?;
+
+    let image = remove_image_border(&image)?;
+
+    image.save_with_format("testing_output.png", ImageFormat::Png)?;
+
+    Ok(())
+  }
+}
@@ -15,252 +15,281 @@ use tokenizers::Tokenizer;
 #[derive(Parser)]
 #[command(author, version, about, long_about = None)]
 struct Args {
-    /// The prompt to be used for image generation.
-    #[arg(long, default_value = "A rusty robot walking on a beach")]
-    prompt: String,
+  /// The prompt to be used for image generation.
+  #[arg(long, default_value = "A rusty robot walking on a beach")]
+  prompt: String,

-    /// Run on CPU rather than on GPU.
-    #[arg(long)]
-    cpu: bool,
+  /// Run on CPU rather than on GPU.
+  #[arg(long)]
+  cpu: bool,

-    /// Use the quantized model.
-    #[arg(long)]
-    quantized: bool,
+  /// Use the quantized model.
+  #[arg(long)]
+  quantized: bool,

-    /// Enable tracing (generates a trace-timestamp.json file).
-    #[arg(long)]
-    tracing: bool,
+  /// Enable tracing (generates a trace-timestamp.json file).
+  #[arg(long)]
+  tracing: bool,

-    /// The height in pixels of the generated image.
-    #[arg(long)]
-    height: Option<usize>,
+  /// The height in pixels of the generated image.
+  #[arg(long)]
+  height: Option<usize>,

-    /// The width in pixels of the generated image.
-    #[arg(long)]
-    width: Option<usize>,
+  /// The width in pixels of the generated image.
+  #[arg(long)]
+  width: Option<usize>,

-    #[arg(long)]
-    decode_only: Option<String>,
+  #[arg(long)]
+  decode_only: Option<String>,

-    #[arg(long, value_enum, default_value = "schnell")]
-    model: Model,
+  #[arg(long, value_enum, default_value = "schnell")]
+  model: Model,

-    /// Use the slower kernels.
-    #[arg(long)]
-    use_dmmv: bool,
+  /// Use the slower kernels.
+  #[arg(long)]
+  use_dmmv: bool,

-    /// The seed to use when generating random samples.
-    #[arg(long)]
-    seed: Option<u64>,
+  /// The seed to use when generating random samples.
+  #[arg(long)]
+  seed: Option<u64>,
+
+  #[arg(long)]
+  cache_path: Option<String>,
+
+  /// Number of blocks to keep active on GPU (default: 1).
+  /// Only relevant when NOT using --use_full_gpu.
+  /// Higher values use more memory but may improve performance.
+  #[arg(long, default_value = "1")]
+  active_blocks: usize,
+
+  /// Use CUDA stream for faster device transfers (default: true).
+  /// Only relevant when running on GPU.
+  #[arg(long, default_value = "true")]
+  use_cuda_stream: bool,
+
+  /// Enable prefetching of the next batch of blocks while processing the current batch.
+  /// Only relevant when NOT using --use_full_gpu.
+  /// This can improve performance by overlapping memory transfers and computation.
+  #[arg(long, default_value = "false")]
+  prefetch_next_batch: bool,
+
+  /// Load entire model on GPU instead of using memory-efficient active blocks approach.
+  /// This uses more GPU memory but avoids CPU-GPU transfers completely.
+  /// When this flag is enabled, the --active_blocks and --prefetch_next_batch flags have no effect.
+  #[arg(long, default_value = "false")]
+  use_full_gpu: bool,
 }

 #[derive(Debug, Clone, Copy, clap::ValueEnum, PartialEq, Eq)]
 enum Model {
-    Schnell,
-    Dev,
+  Schnell,
+  Dev,
 }

 fn run(args: Args) -> Result<()> {
-    use tracing_chrome::ChromeLayerBuilder;
-    use tracing_subscriber::prelude::*;
+  use tracing_chrome::ChromeLayerBuilder;
+  use tracing_subscriber::prelude::*;

-    let Args {
-        prompt,
-        cpu,
-        height,
-        width,
-        tracing,
-        decode_only,
-        model,
-        quantized,
-        ..
-    } = args;
-    let width = width.unwrap_or(1360);
-    let height = height.unwrap_or(768);
+  let Args { prompt, cpu, height, width, tracing, decode_only, model, quantized, active_blocks, use_cuda_stream, prefetch_next_batch, use_full_gpu, .. } = args;
+  let width = width.unwrap_or(1360);
+  let height = height.unwrap_or(768);

-    let _guard = if tracing {
-        let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
-        tracing_subscriber::registry().with(chrome_layer).init();
-        Some(guard)
-    } else {
-        None
+  let _guard = if tracing {
+    let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
+    tracing_subscriber::registry().with(chrome_layer).init();
+    Some(guard)
+  } else {
+    None
+  };
+
+  let api = match args.cache_path.as_ref() {
+    Some(path) => hf_hub::api::sync::ApiBuilder::from_cache(hf_hub::Cache::new(path.to_string().into())).build().map_err(anyhow::Error::msg)?,
+    None => {
+      let cache = hf_hub::Cache::from_env();
+      hf_hub::api::sync::ApiBuilder::from_cache(cache).build().map_err(anyhow::Error::msg)?
+    },
+  };
+  let bf_repo = {
+    let name = match model {
+      Model::Dev => "black-forest-labs/FLUX.1-dev",
+      Model::Schnell => "black-forest-labs/FLUX.1-schnell",
    };
+    api.repo(hf_hub::Repo::model(name.to_string()))
+  };

-    let api = hf_hub::api::sync::Api::new()?;
-    let bf_repo = {
-        let name = match model {
-            Model::Dev => "black-forest-labs/FLUX.1-dev",
-            Model::Schnell => "black-forest-labs/FLUX.1-schnell",
-        };
-        api.repo(hf_hub::Repo::model(name.to_string()))
-    };
-    let device = candle_examples::device(cpu)?;
-    if let Some(seed) = args.seed {
-        device.set_seed(seed)?;
+  // Use CUDA with stream for faster device transfers when not in CPU mode
+  let device = if cpu {
+    candle::Device::Cpu
+  } else {
+    #[cfg(feature = "cuda")]
+    {
+      if use_cuda_stream {
+        println!("Using CUDA with dedicated stream for faster device transfers");
+        // Use device ordinal 0 (first GPU)
+        candle::Device::new_cuda_with_stream(0)?
+      } else {
+        println!("Using standard CUDA device (no dedicated stream)");
+        candle::Device::new_cuda(0)?
+      }
    }
-    let dtype = device.bf16_default_to_f32();
-    let img = match decode_only {
-        None => {
-            let t5_emb = {
-                let repo = api.repo(hf_hub::Repo::with_revision(
-                    "google/t5-v1_1-xxl".to_string(),
-                    hf_hub::RepoType::Model,
-                    "refs/pr/2".to_string(),
-                ));
-                let model_file = repo.get("model.safetensors")?;
-                let vb =
-                    unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
-                let config_filename = repo.get("config.json")?;
-                let config = std::fs::read_to_string(config_filename)?;
-                let config: t5::Config = serde_json::from_str(&config)?;
-                let mut model = t5::T5EncoderModel::load(vb, &config)?;
-                let tokenizer_filename = api
-                    .model("lmz/mt5-tokenizers".to_string())
-                    .get("t5-v1_1-xxl.tokenizer.json")?;
-                let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
-                let mut tokens = tokenizer
-                    .encode(prompt.as_str(), true)
-                    .map_err(E::msg)?
-                    .get_ids()
-                    .to_vec();
-                tokens.resize(256, 0);
-                let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
-                println!("{input_token_ids}");
-                model.forward(&input_token_ids)?
-            };
-            println!("T5\n{t5_emb}");
-            let clip_emb = {
-                let repo = api.repo(hf_hub::Repo::model(
-                    "openai/clip-vit-large-patch14".to_string(),
-                ));
-                let model_file = repo.get("model.safetensors")?;
-                let vb =
-                    unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
-                // https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json
-                let config = clip::text_model::ClipTextConfig {
-                    vocab_size: 49408,
-                    projection_dim: 768,
-                    activation: clip::text_model::Activation::QuickGelu,
-                    intermediate_size: 3072,
-                    embed_dim: 768,
-                    max_position_embeddings: 77,
-                    pad_with: None,
-                    num_hidden_layers: 12,
-                    num_attention_heads: 12,
-                };
-                let model =
-                    clip::text_model::ClipTextTransformer::new(vb.pp("text_model"), &config)?;
-                let tokenizer_filename = repo.get("tokenizer.json")?;
-                let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
-                let tokens = tokenizer
-                    .encode(prompt.as_str(), true)
-                    .map_err(E::msg)?
-                    .get_ids()
-                    .to_vec();
-                let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
-                println!("{input_token_ids}");
-                model.forward(&input_token_ids)?
-            };
-            println!("CLIP\n{clip_emb}");
-            let img = {
-                let cfg = match model {
-                    Model::Dev => flux::model::Config::dev(),
-                    Model::Schnell => flux::model::Config::schnell(),
-                };
-                let img = flux::sampling::get_noise(1, height, width, &device)?.to_dtype(dtype)?;
-                let state = if quantized {
-                    flux::sampling::State::new(
-                        &t5_emb.to_dtype(candle::DType::F32)?,
-                        &clip_emb.to_dtype(candle::DType::F32)?,
-                        &img.to_dtype(candle::DType::F32)?,
-                    )?
-                } else {
-                    flux::sampling::State::new(&t5_emb, &clip_emb, &img)?
-                };
-                let timesteps = match model {
-                    Model::Dev => {
-                        flux::sampling::get_schedule(50, Some((state.img.dim(1)?, 0.5, 1.15)))
-                    }
-                    Model::Schnell => flux::sampling::get_schedule(4, None),
-                };
-                println!("{state:?}");
-                println!("{timesteps:?}");
-                if quantized {
-                    let model_file = match model {
-                        Model::Schnell => api
-                            .repo(hf_hub::Repo::model("lmz/candle-flux".to_string()))
-                            .get("flux1-schnell.gguf")?,
-                        Model::Dev => todo!(),
-                    };
-                    let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(
-                        model_file, &device,
-                    )?;
+    #[cfg(not(feature = "cuda"))]
+    {
+      println!("CUDA not available, using CPU");
+      candle::Device::Cpu
+    }
+  };

-                    let model = flux::quantized_model::Flux::new(&cfg, vb)?;
-                    flux::sampling::denoise(
-                        &model,
-                        &state.img,
-                        &state.img_ids,
-                        &state.txt,
-                        &state.txt_ids,
-                        &state.vec,
-                        &timesteps,
-                        4.,
-                    )?
-                    .to_dtype(dtype)?
-                } else {
-                    let model_file = match model {
-                        Model::Schnell => bf_repo.get("flux1-schnell.safetensors")?,
-                        Model::Dev => bf_repo.get("flux1-dev.safetensors")?,
-                    };
-                    let vb = unsafe {
-                        VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)?
-                    };
-                    let model = flux::model::Flux::new(&cfg, vb)?;
-                    flux::sampling::denoise(
-                        &model,
-                        &state.img,
-                        &state.img_ids,
-                        &state.txt,
-                        &state.txt_ids,
-                        &state.vec,
-                        &timesteps,
-                        4.,
-                    )?
-                }
-            };
-            flux::sampling::unpack(&img, height, width)?
-        }
-        Some(file) => {
-            let mut st = candle::safetensors::load(file, &device)?;
-            st.remove("img").unwrap().to_dtype(dtype)?
-        }
-    };
-    println!("latent img\n{img}");
-
-    let img = {
-        let model_file = bf_repo.get("ae.safetensors")?;
+  if let Some(seed) = args.seed {
+    device.set_seed(seed)?;
+  }
+  let start_time = std::time::Instant::now();
+  let dtype = device.bf16_default_to_f32();
+  let img = match decode_only {
+    None => {
+      let t5_emb = {
+        let repo = api.repo(hf_hub::Repo::with_revision("google/t5-v1_1-xxl".to_string(), hf_hub::RepoType::Model, "refs/pr/2".to_string()));
+        let model_file = repo.get("model.safetensors")?;
        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
+        let config_filename = repo.get("config.json")?;
+        let config = std::fs::read_to_string(config_filename)?;
+        let config: t5::Config = serde_json::from_str(&config)?;
+        let mut model = t5::T5EncoderModel::load(vb, &config)?;
+        let tokenizer_filename = api.model("lmz/mt5-tokenizers".to_string()).get("t5-v1_1-xxl.tokenizer.json")?;
+        let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
+        let mut tokens = tokenizer.encode(prompt.as_str(), true).map_err(E::msg)?.get_ids().to_vec();
+        tokens.resize(256, 0);
+        let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
+        println!("{input_token_ids}");
+        model.forward(&input_token_ids)?
+      };
+      println!("T5\n{t5_emb}");
+      let clip_emb = {
+        let repo = api.repo(hf_hub::Repo::model("openai/clip-vit-large-patch14".to_string()));
+        let model_file = repo.get("model.safetensors")?;
+        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
+        // https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json
+        let config = clip::text_model::ClipTextConfig { vocab_size: 49408, projection_dim: 768, activation: clip::text_model::Activation::QuickGelu, intermediate_size: 3072, embed_dim: 768, max_position_embeddings: 77, pad_with: None, num_hidden_layers: 12, num_attention_heads: 12 };
+        let model = clip::text_model::ClipTextTransformer::new(vb.pp("text_model"), &config)?;
+        let tokenizer_filename = repo.get("tokenizer.json")?;
+        let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
+        let tokens = tokenizer.encode(prompt.as_str(), true).map_err(E::msg)?.get_ids().to_vec();
+        let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
+        println!("{input_token_ids}");
+        model.forward(&input_token_ids)?
+      };
+      println!("CLIP\n{clip_emb}");
+      let img = {
        let cfg = match model {
-            Model::Dev => flux::autoencoder::Config::dev(),
-            Model::Schnell => flux::autoencoder::Config::schnell(),
+          Model::Dev => flux::model::Config::dev(),
+          Model::Schnell => flux::model::Config::schnell(),
        };
-        let model = flux::autoencoder::AutoEncoder::new(&cfg, vb)?;
-        model.decode(&img)?
+        let img = flux::sampling::get_noise(1, height, width, &device)?.to_dtype(dtype)?;
+        let state = if quantized { flux::sampling::State::new(&t5_emb.to_dtype(candle::DType::F32)?, &clip_emb.to_dtype(candle::DType::F32)?, &img.to_dtype(candle::DType::F32)?)? } else { flux::sampling::State::new(&t5_emb, &clip_emb, &img)? };
+        let timesteps = match model {
+          Model::Dev => flux::sampling::get_schedule(28, Some((state.img.dim(1)?, 0.5, 1.15))),
+          Model::Schnell => flux::sampling::get_schedule(4, None),
+        };
+        println!("{state:?}");
+        println!("{timesteps:?}");
+        if quantized {
+          // Note about quantized model memory management
+          let model_file = match model {
+            Model::Schnell => api.repo(hf_hub::Repo::model("lmz/candle-flux".to_string())).get("flux1-schnell.gguf")?,
+            Model::Dev => todo!(),
+          };
+          let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(model_file, &device)?;
+
+          let mut model = flux::quantized_model::Flux::new(&cfg, vb)?;
+          // Quantized models are always fully loaded to the selected device
+          if !cpu {
+            println!("Quantized model fully loaded to GPU with optimized memory layout");
+          }
+
+          flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, &timesteps, 4.)?.to_dtype(dtype)?
+        } else {
+          let model_file = match model {
+            Model::Schnell => bf_repo.get("flux1-schnell.safetensors")?,
+            Model::Dev => bf_repo.get("flux1-dev.safetensors")?,
+          };
+
+          let mut model = if use_full_gpu {
+            // When use_full_gpu is true, load everything directly to GPU
+            println!("Loading entire model directly to GPU...");
+            let gpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
+            let mut model = flux::model::Flux::new(&cfg, gpu_vb)?;
+            model.set_use_device_management(false);
+            println!("Model loaded entirely on GPU");
+            model
+          } else {
+            // Use memory-efficient approach with blocks on CPU
+            // Create a CPU device for initially loading all model weights
+            let cpu_device = candle::Device::Cpu;
+
+            // Load the entire model to CPU first (uses less memory)
+            println!("Loading model to CPU to minimize GPU memory usage...");
+            let cpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &cpu_device)? };
+
+            // Create the model with CPU weights, then selectively move core components to GPU
+            let mut model = flux::model::Flux::new_with_gpu_core(&cfg, cpu_vb, &device)?;
+            println!("Model loaded (core components on GPU, blocks on CPU for reduced memory usage)");
+
+            // Set the number of active blocks to keep on GPU
+            model.set_max_active_blocks(active_blocks);
+            println!("Using {} active blocks on GPU", active_blocks);
+
+            // Enable prefetching if requested
+            if prefetch_next_batch {
+              model.set_prefetch_next_batch(true);
+              println!("Prefetching enabled - will overlap transfers with computation");
+            }
+
+            model
+          };
+
+          // Print final memory management approach being used
+          if use_full_gpu {
+            println!("Using full GPU memory - entire model on GPU");
+          } else {
+            println!("Using device management - blocks will be transferred between CPU and GPU as needed");
+          }
+
+          flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, &timesteps, 4.)?
+        }
+      };
+      flux::sampling::unpack(&img, height, width)?
+    },
+    Some(file) => {
+      let mut st = candle::safetensors::load(file, &device)?;
+      st.remove("img").unwrap().to_dtype(dtype)?
+    },
+  };
+  println!("latent img\n{img}");
+
+  let end_time = std::time::Instant::now();
+  println!("Time taken to decode: {:?}", end_time.duration_since(start_time));
+  let img = {
+    let model_file = bf_repo.get("ae.safetensors")?;
+    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
+    let cfg = match model {
+      Model::Dev => flux::autoencoder::Config::dev(),
+      Model::Schnell => flux::autoencoder::Config::schnell(),
    };
-    println!("img\n{img}");
-    let img = ((img.clamp(-1f32, 1f32)? + 1.0)? * 127.5)?.to_dtype(candle::DType::U8)?;
-    let filename = match args.seed {
-        None => "out.jpg".to_string(),
-        Some(s) => format!("out-{s}.jpg"),
-    };
-    candle_examples::save_image(&img.i(0)?, filename)?;
-    Ok(())
+    let model = flux::autoencoder::AutoEncoder::new(&cfg, vb)?;
+    model.decode(&img)?
+  };
+  println!("img\n{img}");
+  let img = ((img.clamp(-1f32, 1f32)? + 1.0)? * 127.5)?.to_dtype(candle::DType::U8)?;
+  let filename = match args.seed {
+    None => "out.jpg".to_string(),
+    Some(s) => format!("out-{s}.jpg"),
+  };
+  candle_examples::save_image(&img.i(0)?, filename)?;
+  Ok(())
 }

 fn main() -> Result<()> {
-    let args = Args::parse();
-    #[cfg(feature = "cuda")]
-    candle::quantized::cuda::set_force_dmmv(args.use_dmmv);
-    run(args)
+  let args = Args::parse();
+  #[cfg(feature = "cuda")]
+  candle::quantized::cuda::set_force_dmmv(args.use_dmmv);
+  run(args)
 }
@@ -28,7 +28,8 @@
 //! ```
 //!
 //! [`Layer Normalization`]: https://arxiv.org/abs/1607.06450
-use candle::{DType, Module, Result, Tensor, D};
+use candle::{DType, Module, Result, Tensor, D, Device};
+use crate::offload::DeviceTransferable;

 #[derive(Debug, Clone, Copy, PartialEq)]
 pub struct LayerNormConfig {
@@ -201,3 +202,20 @@ pub fn rms_norm(size: usize, eps: f64, vb: crate::VarBuilder) -> Result<RmsNorm>
    };
    Ok(RmsNorm(layer_norm(size, config, vb)?))
 }
+
+impl DeviceTransferable for LayerNorm {
+    fn to_device(&self, device: &Device) -> Result<Self> {
+        Ok(Self {
+            weight: self.weight.to_device(device)?,
+            bias: self.bias.as_ref().map(|b| b.to_device(device)).transpose()?,
+            remove_mean: self.remove_mean,
+            eps: self.eps,
+        })
+    }
+}
+
+impl DeviceTransferable for RmsNorm {
+    fn to_device(&self, device: &Device) -> Result<Self> {
+        Ok(Self(self.0.to_device(device)?))
+    }
+}
@@ -27,6 +27,7 @@ pub mod kv_cache;
 pub mod layer_norm;
 pub mod linear;
 pub mod loss;
+pub mod offload;
 pub mod ops;
 pub mod optim;
 pub mod rnn;
@@ -58,3 +59,4 @@ pub use var_builder::VarBuilder;
 pub use var_map::VarMap;

 pub use candle::{Module, ModuleT};
+pub use offload::DeviceTransferable;
@@ -17,7 +17,9 @@
 //! assert_eq!(ys.to_vec2::<f32>()?, &[[210.0, 430.0, 650.0]]);
 //! # Ok(()) }
 //! ```
-use candle::{Result, Tensor};
+use candle::{Result, Tensor, Device};
+use crate::offload::DeviceTransferable;
+

 #[derive(Clone, Debug)]
 pub struct Linear {
@@ -88,3 +90,12 @@ pub fn linear_b(
        linear_no_bias(in_dim, out_dim, vb)
    }
 }
+
+impl DeviceTransferable for Linear {
+    fn to_device(&self, device: &Device) -> Result<Self> {
+        Ok(Self {
+            weight: self.weight.to_device(device)?,
+            bias: self.bias.as_ref().map(|b| b.to_device(device)).transpose()?,
+        })
+    }
+}
@@ -0,0 +1,18 @@
+use candle::{Device, Result};
+
+pub trait DeviceTransferable {
+    /// Move the module to the specified device
+    ///
+    /// # Arguments
+    ///
+    /// * `device` - The target device to move the module to
+    ///
+    /// # Returns
+    ///
+    /// A new instance of the module with all tensors moved to the target device,
+    /// or an error if the module cannot be moved.
+    fn to_device(&self, _device: &Device) -> Result<Self> where Self: Sized {
+        // Default implementation that rejects moving - components can override this
+        candle::bail!("to_device not implemented for this type")
+    }
+}
@@ -9,7 +9,9 @@ candle = { path = "../candle-core", package = "candle-core" }
 candle-flash-attn = { path = "../candle-flash-attn", optional = true }
 candle-nn = { path = "../candle-nn", package = "candle-nn" }
 fancy-regex = "0.13.0"
+half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
 intel-mkl-src = { version = "0.8.1", features = ["mkl-static-lp64-iomp"], optional = true }
+log = "0.4.27"
 num-traits = "0.2.15"
 rand = "0.8.5"
 rayon = "1.7.0"
@@ -17,7 +19,6 @@ serde = { version = "1.0.171", features = ["derive"] }
 serde_json = "1.0.99"
 serde_plain = "1.0.2"
 tracing = "0.1.37"
-half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }

 [features]
 default = []
@@ -24,9 +24,39 @@
 use candle::{Result, Tensor};

 pub trait WithForward {
+    /// The forward pass for the Flux model.
+    ///
+    /// # Arguments
+    ///
+    /// * `img` - The input image tensor
+    /// * `img_ids` - The image position ids
+    /// * `txt` - The text tensor
+    /// * `txt_ids` - The text position ids
+    /// * `timesteps` - The timesteps tensor
+    /// * `y` - The conditioning vector
+    /// * `guidance` - An optional guidance vector
+    ///
+    /// # Returns
+    ///
+    /// The result of the forward pass
+    ///
+    /// # Note on Memory Optimization
+    ///
+    /// To optimize memory usage, we could implement a mechanism to move model blocks 
+    /// to CPU when not in use. This would require:
+    ///
+    /// 1. Adding `to_device` methods to all model blocks (DoubleStreamBlock, SingleStreamBlock, etc.)
+    /// 2. Implementing these methods by moving all tensors inside each block to the target device
+    /// 3. Modifying the forward method to:
+    ///    - Move each block to GPU when needed
+    ///    - Process the block
+    ///    - Move the block back to CPU when done
+    ///
+    /// This pattern could significantly reduce GPU memory usage for large models by
+    /// keeping only the currently active block in GPU memory.
    #[allow(clippy::too_many_arguments)]
    fn forward(
-        &self,
+        &mut self,
        img: &Tensor,
        img_ids: &Tensor,
        txt: &Tensor,
@@ -419,7 +419,7 @@ impl Flux {
 impl super::WithForward for Flux {
    #[allow(clippy::too_many_arguments)]
    fn forward(
-        &self,
+        &mut self,
        img: &Tensor,
        img_ids: &Tensor,
        txt: &Tensor,
@@ -93,7 +93,7 @@ pub fn unpack(xs: &Tensor, height: usize, width: usize) -> Result<Tensor> {

 #[allow(clippy::too_many_arguments)]
 pub fn denoise<M: super::WithForward>(
-    model: &M,
+    model: &mut M,
    img: &Tensor,
    img_ids: &Tensor,
    txt: &Tensor,
@@ -1,6 +1,9 @@
 use super::schedulers::{betas_for_alpha_bar, BetaSchedule, PredictionType};
 use candle::{DType, Device, IndexOp, Result, Tensor};
+use log::debug;
+
 const FINAL_ALPHA_CUMPROD: f64 = 0.9991499781608582;
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum LCMVarianceType {
    FixedSmall,
@@ -87,7 +90,7 @@ impl LCMScheduler {
        let timesteps = LCMScheduler::get_timesteps_for_steps(inference_steps, strength);
        let step_ratio = config.train_timesteps / inference_steps;

-        println!("Using LCM timesteps: {:?}", timesteps);
+        debug!("Using LCM timesteps: {:?}", timesteps);

        Ok(Self {
            alphas_cumprod,
@@ -143,11 +146,11 @@ impl LCMScheduler {
        let scaled_timestep_squared = scaled_timestep.powi(2);
        let denominator = scaled_timestep_squared + sigma_squared;

-        println!(
+        debug!(
            "Boundary condition for timestep {}, scaled timestep {}",
            timestep, scaled_timestep
        );
-        println!(
+        debug!(
            "  Intermediate values: sigma_data^2: {:.20}, scaled_timestep^2: {:.20}, denominator: {:.20}",
            sigma_squared, scaled_timestep_squared, denominator
        );
@@ -155,7 +158,7 @@ impl LCMScheduler {
        let c_skip = sigma_squared / denominator;
        let c_out = scaled_timestep / (scaled_timestep_squared + sigma_squared).sqrt();

-        println!(
+        debug!(
            "  Final computed values: c_skip: {:.20}, c_out: {:.20}",
            c_skip, c_out
        );
@@ -172,14 +175,14 @@ impl LCMScheduler {
    ) -> Result<Tensor> {
        // Scale by exactly 1000.0 to match diffusers
        let w = guidance_scale * 1000.0;
-        println!(
+        debug!(
            "Creating guidance embedding with dim={}, scaled value={}",
            embedding_dim, w
        );

        // Calculate the half dimension
        let half_dim = embedding_dim / 2;
-        println!("Half dimension: {}", half_dim);
+        debug!("Half dimension: {}", half_dim);

        let log_10000 = (10000.0f32).ln();
        let freq_factor = log_10000 / ((half_dim - 1) as f32); // pre-compute division
@@ -191,7 +194,7 @@ impl LCMScheduler {
            freqs.push(freq);
        }

-        println!(
+        debug!(
            "First few frequencies: {:?}",
            &freqs.iter().take(5).collect::<Vec<_>>()
        );
@@ -204,39 +207,39 @@ impl LCMScheduler {
            emb_data.push(w_val * freq);
        }

-        println!(
+        debug!(
            "First few scaled frequencies: {:?}",
            &emb_data.iter().take(5).collect::<Vec<_>>()
        );

        // Create the embedding tensor
        let emb = Tensor::new(emb_data.as_slice(), device)?.to_dtype(dtype)?;
-        println!("Embedding tensor shape: {:?}", emb.shape());
+        debug!("Embedding tensor shape: {:?}", emb.shape());

        // Reshape for correct sin/cos application
        let emb = emb.reshape((1, half_dim))?;
-        println!("Reshaped embedding tensor: {:?}", emb.shape());
+        debug!("Reshaped embedding tensor: {:?}", emb.shape());

        // Compute sin and cos components
        let sin_emb = emb.sin()?;
-        println!("Sin embedding shape: {:?}", sin_emb.shape());
+        debug!("Sin embedding shape: {:?}", sin_emb.shape());
        let cos_emb = emb.cos()?;
-        println!("Cos embedding shape: {:?}", cos_emb.shape());
+        debug!("Cos embedding shape: {:?}", cos_emb.shape());

        // Concatenate sin and cos (along dimension 1)
        let result = Tensor::cat(&[sin_emb, cos_emb], 1)?;
-        println!("After concatenation shape: {:?}", result.shape());
+        debug!("After concatenation shape: {:?}", result.shape());

        // If odd dimension, pad with zero
        let result = if embedding_dim % 2 == 1 {
-            println!("Adding zero padding for odd dimension");
+            debug!("Adding zero padding for odd dimension");
            let zero = Tensor::zeros((1, 1), dtype, device)?;
            Tensor::cat(&[result, zero], 1)?
        } else {
            result
        };

-        println!("Final guidance embedding shape: {:?}", result.shape());
+        debug!("Final guidance embedding shape: {:?}", result.shape());

        Ok(result)
    }
@@ -288,12 +291,12 @@ impl LCMScheduler {
                };

                let timesteps = linspace.iter().map(|&f| f as usize).rev().collect();
-                println!("timesteps: {:?}", timesteps);
+                debug!("timesteps: {:?}", timesteps);
                timesteps
            }
        };

-        println!("all_timesteps: {:?}", all_timesteps);
+        debug!("all_timesteps: {:?}", all_timesteps);

        // Compute `t_start` based on `strength`
        let num_train_timesteps = 1000;
@@ -303,7 +306,7 @@ impl LCMScheduler {
            .position(|&t| t <= init_timestep)
            .unwrap_or(0);

-        println!(
+        debug!(
            "get_timesteps: t_start: {} init_timestep: {} strength: {}",
            t_start, init_timestep, strength
        );
@@ -327,7 +330,7 @@ impl LCMScheduler {
            0
        };

-        println!(
+        debug!(
            "LCM step: timestep={}, prev_timestep={}, step_index={}",
            timestep, prev_timestep, step_index
        );
@@ -357,15 +360,15 @@ impl LCMScheduler {
        let denoised = ((c_out * pred_original_sample )? + (c_skip * sample)?)?;

        let pred_prev_sample = if step_index != self.timesteps.len() - 1 {
-            println!("Adding noise at intermediate step {}", timestep);
+            debug!("Adding noise at intermediate step {}", timestep);
            let noise = Tensor::randn_like(&denoised, 0.0, 1.0)?;
            ((alpha_prod_t_prev.sqrt() * denoised)? + (beta_prod_t_prev.sqrt() * noise)?)?
        } else {
-            println!("No noise for final timestep {}", timestep);
+            debug!("No noise for final timestep {}", timestep);
            denoised
        };

-        println!(
+        debug!(
            "Denoised sample shape: {:?}, first few values: {}",
            pred_prev_sample.shape(),
            pred_prev_sample
@@ -388,7 +391,7 @@ impl LCMScheduler {
    //         self.timesteps[step_index + 1]
    //     };

-    //     println!(
+    //     debug!(
    //         "LCM step: timestep={}, prev_timestep={}",
    //         timestep, prev_timestep
    //     );
@@ -401,7 +404,7 @@ impl LCMScheduler {
    //         self.alphas_cumprod[prev_timestep]
    //     };

-    //     println!(
+    //     debug!(
    //         "LCM scheduler step: timestep={}, step_index={}, sigma={}",
    //         timestep,
    //         step_index,
@@ -433,13 +436,13 @@ impl LCMScheduler {
    //         pred_original_sample = pred_original_sample.clamp(-1f32, 1f32)?;
    //     }

-    //     println!("boundary: Pred original sample first few values: {}",
+    //     debug!("boundary: Pred original sample first few values: {}",
    //              pred_original_sample.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);

    //     // Denoise model output using boundary conditions
    //     let denoised = ((pred_original_sample * c_out)? + (sample * c_skip)?)?;

-    //     println!("boundary: Denoised first few values: {}",
+    //     debug!("boundary: Denoised first few values: {}",
    //              denoised.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);

    //     // LCM doesn't add noise during inference
@@ -452,11 +455,11 @@ impl LCMScheduler {

    //     // let pred_prev_sample = pred_prev_sample.clamp(-1f32, 1f32)?;

-    //     println!(
+    //     debug!(
    //         "LCM scheduler pred_original shape: {:?}",
    //         pred_prev_sample.shape()
    //     );
-    //     println!(
+    //     debug!(
    //         "LCM scheduler pred_original first few values: {}",
    //         pred_prev_sample
    //             .i(0)?
@@ -469,7 +472,7 @@ impl LCMScheduler {
    // }

    pub fn add_noise(&self, original: &Tensor, noise: Tensor, timestep: usize) -> Result<Tensor> {
-        println!("Adding noise at timestep {}", timestep);
+        debug!("Adding noise at timestep {}", timestep);
        let timestep = if timestep >= self.alphas_cumprod.len() {
            self.alphas_cumprod.len() - 1 // Ensure valid range
        } else {
@@ -491,7 +494,7 @@ impl LCMScheduler {
    //     let alpha_cumprod = self.alphas_cumprod[timestep];

    //     // Print the alpha value to debug
-    //     println!(
+    //     debug!(
    //         "Adding noise with timestep={}, alpha={}, sqrt(alpha)={}",
    //         timestep,
    //         alpha_cumprod,
@@ -0,0 +1,26 @@
+# Logs
+logs
+*.log
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+pnpm-debug.log*
+lerna-debug.log*
+
+node_modules
+dist
+dist-ssr
+*.local
+backup
+.venv
+# Editor directories and files
+~
+.vscode
+.idea
+.DS_Store
+*.suo
+*.ntvs*
+*.njsproj
+*.sln
+*.sw?
+**/*.onnx
@@ -0,0 +1,8 @@
+<html>
+  <head>
+    <title>Not Found</title>
+  </head>
+  <body>
+    <h1> This Page Was Not Found </h1>
+  </body>
+</html>
@@ -0,0 +1,39 @@
+# Memeboard
+
+## Installation & Dev
+
+```sh
+npm install
+npm run dev
+```
+
+or
+
+```sh
+yarn      # this installs the node packages
+yarn dev  # this runs the vite dev server
+```
+
+You may need to turn off browser security to get around CORS restrictions:
+
+### Windows 11
+
+from your powershell
+
+```sh
+# Chrome
+Start "C:\Program Files\Google\Chrome\Application\chrome.exe" -ArgumentList '--user-data-dir="C://chrome-dev-disabled-security" --disable-web-security --disable-site-isolation-trials'
+```
+
+### on MAC
+
+form your terminal
+
+```sh
+# Chrome
+/Applications/Google\ Chrome.app/Contents/MacOS/Google\ Chrome --disable-web-security --user-data-dir=~/chrome
+```
+
+### Inspecting the Shared Workers
+
+on chrome, type `chrome://inspect/#workers` into the url field
@@ -0,0 +1,3 @@
+/*
+  Access-Control-Allow-Origin: *
+
@@ -0,0 +1,11 @@
+# Board Netlify Redirects File
+
+# Netlify docs: 
+#   The redirects engine will process the first matching rule it 
+#   finds, reading from top to bottom. Rules in the _redirects file 
+#   are always processed first, followed by rules in the Netlify 
+#   configuration file.
+
+/landing               /index.html                   200!
+/login                 /index.html                   200!
+/m_*                   /index.html                   200!
@@ -0,0 +1,44 @@
+import js from "@eslint/js";
+import globals from "globals";
+import reactHooks from "eslint-plugin-react-hooks";
+import reactRefresh from "eslint-plugin-react-refresh";
+import tseslint from "typescript-eslint";
+import react from "eslint-plugin-react";
+
+export default tseslint.config(
+  { ignores: ["dist"] },
+  {
+    extends: [
+      js.configs.recommended,
+      ...tseslint.configs.recommendedTypeChecked,
+      ...tseslint.configs.stylisticTypeChecked,
+    ],
+    files: ["**/*.{ts,tsx}"],
+    languageOptions: {
+      ecmaVersion: 2020,
+      globals: globals.browser,
+    },
+    plugins: {
+      "react-hooks": reactHooks,
+      "react-refresh": reactRefresh,
+      react,
+    },
+    rules: {
+      ...reactHooks.configs.recommended.rules,
+      "react-refresh/only-export-components": [
+        "warn",
+        { allowConstantExport: true },
+      ],
+      ...react.configs.recommended.rules,
+      ...react.configs["jsx-runtime"].rules,
+      // TODO WIL: make no-reaching rule for imports
+      // "import/no-internal-modules": [
+      //   "error",
+      //   {
+      //     allow: [],
+      //     forbid: ["**/uiAccess/*"],
+      //   },
+      // ],
+    },
+  },
+);
@@ -0,0 +1,24 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <link rel="icon" type="image/x-ico" href="/favicon.ico" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <link rel="preconnect" href="https://fonts.googleapis.com" />
+    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin />
+    <link
+      href="https://fonts.googleapis.com/css2?family=Source+Sans+3:ital,wght@0,200..900;1,200..900&display=swap"
+      rel="stylesheet"
+    />
+    <link
+      href="https://fonts.googleapis.com/css2?family=Syne:wght@600&display=swap"
+      rel="stylesheet"
+    />
+    <title>Board | Storyteller.AI</title>
+  </head>
+
+  <body>
+    <div id="root"></div>
+    <script type="module" src="/src/app/index.tsx"></script>
+  </body>
+</html>
@@ -0,0 +1,34 @@
+[build]
+command = "./script/netlify_build.sh"
+publish = "dist/"
+
+#[build.environment]
+#REACT_APP_PUBLIC_POSTHOG_KEY = "phc_jBFgac0mVALAFk3negnSfYCcHgvkT00yBLQDmCDYNBb"
+#REACT_APP_PUBLIC_POSTHOG_UI = "https://us.i.posthog.com"
+
+#[dev]
+#command = "npm run dev"
+#framework = "vite"
+
+# Set immutable caching for static files, because they have fingerprinted filenames
+
+[[headers]]
+for = "/build/*"
+
+[headers.values]
+
+"Cache-Control" = "public, max-age=31560000, immutable"
+
+# [[redirects]]
+#   from = "/ingest/static/*"
+#   to = "https://us-assets.i.posthog.com/static/:splat"
+#   host = "us-assets.i.posthog.com"
+#   status = 200
+#   force = true
+# 
+# [[redirects]]
+#   from = "/ingest/*"
+#   to = "https://us.i.posthog.com/:splat"
+#   host = "us.i.posthog.com"
+#   status = 200
+#   force = true
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Michael Chung	1ce9c8bed8	Added reference images to the generation enqueueImageGeneration now accepts additionalImages enqueueImageGeneration needs testing. Polling for jobs doesn't work, most api calls don't work and need integration testing.	2025-04-03 17:26:52 -04:00
Michael Chung	25c15bffd8	Merge branch 'main' into feature/gen	2025-04-03 12:37:31 -04:00
bflatastic	8ff91690c1	bring back objects, skybox, creatures tabs	2025-04-03 23:24:14 +07:00
BFlat	97bbe0fe07	add image ui (#471 )	2025-04-03 23:01:02 +07:00
Michael Chung	2d361c6636	Merge branch 'main' into feature/gen	2025-04-03 11:55:38 -04:00
BFlat	9361634d83	final fix ratio default to 3:2 (#469 )	2025-04-03 12:41:40 +07:00
bflatastic	7cf0363555	fixed ratio in render 3:2	2025-04-03 12:32:45 +07:00
bflatastic	c2ee3b0ce1	render default to 3:2 ratio	2025-04-03 12:30:57 +07:00
BFlat	ca6fc13171	aspect ratio update (#468 ) - added 3:2, 2:3 ratios (defaulted to 3:2 [1200 x 800 px]) - added letterbox accordingly - aspect ratio buttons to show icons for each ratio - popover to close when item selected	2025-04-03 12:23:54 +07:00
Michael Chung	534afbbbf9	Wire screenshot upload and generate to the engine.	2025-04-02 21:36:42 -04:00
Brandon Thomas	cb1ea1f828	fix CDN in local development	2025-04-02 18:48:44 -04:00
bflatastic	010e3be5dc	make letterbox always visible	2025-04-02 22:59:59 +07:00
bflatastic	d8dab1a66e	fix trimming 2	2025-04-02 22:42:28 +07:00
bflatastic	97b236d04d	paste trim fix	2025-04-02 22:40:56 +07:00
bflatastic	15fe59cc19	added enqueueing state to generate button, limited textarea height and updates	2025-04-02 22:36:53 +07:00
BFlat	89bec31540	removed video controls and add prompt box (#467 )	2025-04-02 22:01:19 +07:00
ArEnScandMichael Chung	7568a0ef86	Fixes for the new 3d editor (#465 ) Removed the login page and timeline and confusing UI components. Extended scene to the bottom of the screen. Co-authored-by: Michael Chung <michaelchung@michaels-mbp-2.home>	2025-04-01 19:11:42 -04:00
Brandon Thomas	6180222e15	netlify.toml for editor2d	2025-04-01 00:48:53 -04:00
Brandon Thomas	6d45ce29e2	Adding copy of `realtime` as `editor2d` for web deployment	2025-04-01 00:44:28 -04:00
Brandon Thomas	1da56d6d22	netlify build outputs dir	2025-04-01 00:06:22 -04:00
Brandon Thomas	91a1c00442	reserved names	2025-04-01 00:00:52 -04:00
Brandon Thomas	3db942f1f6	artcraft cors	2025-03-31 23:58:48 -04:00
Brandon Thomas	2e56291649	unblock rust docker builds	2025-03-31 21:41:54 -04:00
Brandon Thomas	b31badc6d1	Editor3D - this directory is a fork of "studio gen 2" by Tanish. We'll be using this for our image generation roadmap: https://docs.google.com/document/d/1O9YI_iLwXkl2wAJJ23dFp3eHR-JzhlRHCBPD8psG-1E/edit The original Studio Gen 2 is in the `frontend/studio-gen-2` directory. The original Studio Gen 1 is in the `frontend/studio-gen-1` directory	2025-03-31 13:28:59 -04:00
Brandon Thomas	ed399fcb21	rename storyteller animator to "studio-gen-2"	2025-03-31 13:25:51 -04:00
Brandon Thomas	d858c9394b	work in progress background removal adjustment	2025-03-28 01:33:52 -04:00
Brandon Thomas	a19a6599ec	rename model	2025-03-27 23:02:51 -04:00
Brandon Thomas	57c15f7c9f	rename realtime to sketch	2025-03-27 22:41:40 -04:00
Brandon Thomas	aa25ae8c87	load flux models from R2	2025-03-27 21:04:39 -04:00
Brandon Thomas	c27055761e	fix	2025-03-26 20:35:50 -04:00
Brandon Thomas	abf3024892	add to downloaders	2025-03-26 20:35:11 -04:00
Brandon Thomas	2182d4dde1	lots of new models in R2	2025-03-26 20:32:05 -04:00
Brandon Thomas	b632b927a3	save images to file	2025-03-26 13:12:54 -04:00
Brandon Thomas	4e2be12064	flux cherry picked in	2025-03-26 02:37:04 -04:00
Brandon Thomas	29b33cd327	save image handler	2025-03-25 22:43:25 -04:00
Brandon Thomas	032d0e1eae	remove lock button	2025-03-25 20:05:49 -04:00
Brandon Thomas	a3ddffeeee	get rid of "transform" / "move" context menu	2025-03-25 19:54:51 -04:00
Brandon Thomas	28a5ff2f9f	remove background from preview images	2025-03-25 19:43:48 -04:00
Brandon Thomas	1028e04e02	turn lcm println!() to debug!() level	2025-03-25 17:45:29 -04:00
Brandon Thomas	49c010db59	temporarily remove references to flux (temp fix)	2025-03-25 17:34:33 -04:00
Kasisnu	48ff60d3b8	super basic flux (#463 )	2025-03-25 17:05:50 -04:00
Brandon Thomas	9c5b993c4f	read seed from config too	2025-03-25 12:34:33 -04:00
Brandon Thomas	a4aecc4d38	send seed as a parameter	2025-03-25 12:21:47 -04:00
Brandon Thomas	4d73230e84	add logging	2025-03-25 10:44:11 -04:00
Brandon Thomas	a5a09f3e64	fixup inpaint stub	2025-03-25 01:44:20 -04:00
Brandon Thomas Tanish Manku bflatastic	49894bd9b3	Generate Mode (#461 ) Co-authored-by: Tanish Manku <-h> Co-authored-by: bflatastic <bombayapirat@gmail.com>	2025-03-25 01:38:56 -04:00
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	e3245dc593	fix on windows	2025-03-22 00:58:04 -04: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
Brandon Thomas	5fdb4d1024	disable msi builds due to 2.0 gb limit for 32bit	2025-03-20 11:59:43 -04: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
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