diff --git a/Cargo.lock b/Cargo.lock index 7932ece0cf..bddcc471b8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -390,6 +390,12 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" +[[package]] +name = "adler2" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627" + [[package]] name = "aes" version = "0.7.5" @@ -597,15 +603,22 @@ checksum = "b3d1d046238990b9cf5bcde22a3fb3584ee5cf65fb2765f454ed428c7a0063da" name = "app" version = "0.1.0" dependencies = [ + "anyhow", "base64 0.22.1", + "candle-core", + "candle-nn", + "candle-transformers", + "hf-hub", "image", "log", "once_cell", + "rand 0.8.5", "serde", "serde_json", "tauri", "tauri-build", "tauri-plugin-log", + "tokenizers", ] [[package]] @@ -930,7 +943,7 @@ dependencies = [ "cc", "cfg-if 1.0.0", "libc", - "miniz_oxide", + "miniz_oxide 0.7.2", "object", "rustc-demangle", ] @@ -1027,12 +1040,38 @@ dependencies = [ "utoipa", ] +[[package]] +name = "bindgen_cuda" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f8489af5b7d17a81bffe37e0f4d6e1e4de87c87329d05447f22c35d95a1227d" +dependencies = [ + "glob", + "num_cpus", + "rayon", +] + [[package]] name = "binstring" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e0d60973d9320722cb1206f412740e162a33b8547ea8d6be75d7cff237c7a85" +[[package]] +name = "bit-set" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" + [[package]] name = "bit_field" version = "0.10.2" @@ -1269,6 +1308,20 @@ name = "bytemuck" version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78834c15cb5d5efe3452d58b1e8ba890dd62d21907f867f383358198e56ebca5" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fa76293b4f7bb636ab88fd78228235b5248b4d05cc589aed610f954af5d7c7a" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.87", +] [[package]] name = "byteorder" @@ -1365,6 +1418,74 @@ dependencies = [ "serde", ] +[[package]] +name = "candle-core" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "855dfedff437d2681d68e1f34ae559d88b0dd84aa5a6b63f2c8e75ebdd875bbf" +dependencies = [ + "byteorder", + "candle-kernels", + "cudarc", + "gemm", + "half", + "memmap2", + "num-traits", + "num_cpus", + "rand 0.8.5", + "rand_distr", + "rayon", + "safetensors", + "thiserror 1.0.63", + "ug", + "ug-cuda", + "yoke", + "zip 1.1.4", +] + +[[package]] +name = "candle-kernels" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53343628fa470b7075c28c589b98735b4220b464e37ddbb8e117040e199f4787" +dependencies = [ + "bindgen_cuda", +] + +[[package]] +name = "candle-nn" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd3c6b2ee0dfd64af12ae5b07e4b7c517898981cdaeffcb10b71d7dd5c8f359" +dependencies = [ + "candle-core", + "half", + "num-traits", + "rayon", + "safetensors", + "serde", + "thiserror 1.0.63", +] + +[[package]] +name = "candle-transformers" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4270cc692c4a3df2051c2e8c3c4da3a189746af7ca3a547b99ecd335582b92e1" +dependencies = [ + "byteorder", + "candle-core", + "candle-nn", + "fancy-regex", + "num-traits", + "rand 0.8.5", + "rayon", + "serde", + "serde_json", + "serde_plain", + "tracing", +] + [[package]] name = "cargo-platform" version = "0.1.9" @@ -1737,6 +1858,19 @@ dependencies = [ "url", ] +[[package]] +name = "console" +version = "0.15.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea3c6ecd8059b57859df5c69830340ed3c41d30e3da0c1cbed90a96ac853041b" +dependencies = [ + "encode_unicode", + "libc", + "once_cell", + "unicode-width 0.2.0", + "windows-sys 0.59.0", +] + [[package]] name = "const-oid" version = "0.9.6" @@ -2088,6 +2222,16 @@ dependencies = [ "syn 2.0.87", ] +[[package]] +name = "cudarc" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38cd60a9a42ec83a2ed7effb0b1f073270264ea99da7acfc44f7e8d74dee0384" +dependencies = [ + "half", + "libloading 0.8.6", +] + [[package]] name = "darling" version = "0.13.4" @@ -2733,6 +2877,16 @@ version = "1.0.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "545b22097d44f8a9581187cdf93de7a71e4722bf51200cfaba810865b49a495d" +[[package]] +name = "dyn-stack" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56e53799688f5632f364f8fb387488dd05db9fe45db7011be066fc20e7027f8b" +dependencies = [ + "bytemuck", + "reborrow", +] + [[package]] name = "easyenv" version = "0.1.0" @@ -2948,6 +3102,12 @@ version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4ef6b89e5b37196644d8796de5268852ff179b44e96276cf4290264843743bb7" +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + [[package]] name = "encoding_rs" version = "0.8.34" @@ -2957,6 +3117,18 @@ dependencies = [ "cfg-if 1.0.0", ] +[[package]] +name = "enum-as-inner" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1e6a265c649f3f5979b601d26f1d05ada116434c87741c9493cb56218f76cbc" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "quote", + "syn 2.0.87", +] + [[package]] name = "enums" version = "0.0.1" @@ -3077,6 +3249,15 @@ dependencies = [ "tokio", ] +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + [[package]] name = "etcetera" version = "0.8.0" @@ -3096,20 +3277,30 @@ checksum = "0206175f82b8d6bf6652ff7d71a1e27fd2e4efde587fd368662814d6ec1d9ce0" [[package]] name = "exr" -version = "1.71.0" +version = "1.73.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "832a761f35ab3e6664babfbdc6cef35a4860e816ec3916dcfd0882954e98a8a8" +checksum = "f83197f59927b46c04a183a619b7c29df34e63e63c7869320862268c0ef687e0" dependencies = [ "bit_field", - "flume", "half", "lebe", - "miniz_oxide", + "miniz_oxide 0.8.4", "rayon-core", "smallvec 1.13.2", "zune-inflate", ] +[[package]] +name = "fancy-regex" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "531e46835a22af56d1e3b66f04844bed63158bc094a628bec1d321d9b4c44bf2" +dependencies = [ + "bit-set", + "regex-automata", + "regex-syntax", +] + [[package]] name = "fastrand" version = "1.9.0" @@ -3197,7 +3388,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5f54427cfd1c7829e2a139fcefea601bf088ebca651d2bf53ebc600eac295dae" dependencies = [ "crc32fast", - "miniz_oxide", + "miniz_oxide 0.7.2", ] [[package]] @@ -3542,6 +3733,124 @@ dependencies = [ "x11", ] +[[package]] +name = "gemm" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ab24cc62135b40090e31a76a9b2766a501979f3070fa27f689c27ec04377d32" +dependencies = [ + "dyn-stack", + "gemm-c32", + "gemm-c64", + "gemm-common", + "gemm-f16", + "gemm-f32", + "gemm-f64", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-c32" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9c030d0b983d1e34a546b86e08f600c11696fde16199f971cd46c12e67512c0" +dependencies = [ + "dyn-stack", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-c64" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbb5f2e79fefb9693d18e1066a557b4546cd334b226beadc68b11a8f9431852a" +dependencies = [ + "dyn-stack", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-common" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2e7ea062c987abcd8db95db917b4ffb4ecdfd0668471d8dc54734fdff2354e8" +dependencies = [ + "bytemuck", + "dyn-stack", + "half", + "num-complex", + "num-traits", + "once_cell", + "paste", + "pulp", + "raw-cpuid", + "rayon", + "seq-macro", + "sysctl", +] + +[[package]] +name = "gemm-f16" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ca4c06b9b11952071d317604acb332e924e817bd891bec8dfb494168c7cedd4" +dependencies = [ + "dyn-stack", + "gemm-common", + "gemm-f32", + "half", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "rayon", + "seq-macro", +] + +[[package]] +name = "gemm-f32" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9a69f51aaefbd9cf12d18faf273d3e982d9d711f60775645ed5c8047b4ae113" +dependencies = [ + "dyn-stack", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-f64" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa397a48544fadf0b81ec8741e5c0fba0043008113f71f2034def1935645d2b0" +dependencies = [ + "dyn-stack", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -3559,7 +3868,7 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "14dbbfd5c71d70241ecf9e6f13737f7b5ce823821063188d7e46c41d371eebd5" dependencies = [ - "unicode-width", + "unicode-width 0.1.11", ] [[package]] @@ -3830,11 +4139,16 @@ dependencies = [ [[package]] name = "half" -version = "2.2.1" +version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02b4af3693f1b705df946e9fe5631932443781d0aabb423b62fcd4d73f6d2fd0" +checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" dependencies = [ + "bytemuck", + "cfg-if 1.0.0", "crunchy", + "num-traits", + "rand 0.8.5", + "rand_distr", ] [[package]] @@ -3904,6 +4218,23 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "hf-hub" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b780635574b3d92f036890d8373433d6f9fc7abb320ee42a5c25897fc8ed732" +dependencies = [ + "dirs 5.0.1", + "indicatif", + "log", + "native-tls", + "rand 0.8.5", + "serde", + "serde_json", + "thiserror 1.0.63", + "ureq", +] + [[package]] name = "hkdf" version = "0.12.3" @@ -4373,6 +4704,19 @@ dependencies = [ "serde", ] +[[package]] +name = "indicatif" +version = "0.17.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" +dependencies = [ + "console", + "number_prefix", + "portable-atomic", + "unicode-width 0.2.0", + "web-time", +] + [[package]] name = "infer" version = "0.2.3" @@ -4847,7 +5191,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e9ec52138abedcc58dc17a7c6c0c00a2bdb4f3427c7f63fa97fd0d859155caf" dependencies = [ "gtk-sys", - "libloading", + "libloading 0.7.4", "once_cell", ] @@ -4877,6 +5221,16 @@ dependencies = [ "winapi 0.3.9", ] +[[package]] +name = "libloading" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc2f4eb4bc735547cfed7c0a4922cbd04a4655978c09b54f1f7b228750664c34" +dependencies = [ + "cfg-if 1.0.0", + "windows-targets 0.52.6", +] + [[package]] name = "libm" version = "0.2.8" @@ -5039,6 +5393,22 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4" +[[package]] +name = "macro_rules_attribute" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a82271f7bc033d84bbca59a3ce3e4159938cb08a9c3aebbe54d215131518a13" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8dd856d451cc0da70e2ef2ce95a18e39a93b7558bedf10201ad28503f918568" + [[package]] name = "magic-crypt" version = "3.1.13" @@ -5184,6 +5554,16 @@ version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +[[package]] +name = "memmap2" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd3f7eed9d3848f8b98834af67102b720745c4ec028fcd0aa0239277e7de374f" +dependencies = [ + "libc", + "stable_deref_trait", +] + [[package]] name = "memoffset" version = "0.7.1" @@ -5312,6 +5692,15 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "miniz_oxide" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3b1c9bd4fe1f0f8b387f6eb9eb3b4a1aa26185e5750efb9140301703f62cd1b" +dependencies = [ + "adler2", +] + [[package]] name = "mio" version = "0.6.23" @@ -5405,6 +5794,27 @@ dependencies = [ "syn 2.0.87", ] +[[package]] +name = "monostate" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d208407d7552cd041d8cdb69a1bc3303e029c598738177a3d87082004dc0e1e" +dependencies = [ + "monostate-impl", + "serde", +] + +[[package]] +name = "monostate-impl" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7ce64b975ed4f123575d11afd9491f2e37bbd5813fbfbc0f09ae1fbddea74e0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.87", +] + [[package]] name = "mp4" version = "0.14.0" @@ -5489,7 +5899,7 @@ dependencies = [ "log", "md-5 0.10.6", "memchr", - "miniz_oxide", + "miniz_oxide 0.7.2", "native-tls", "num-bigint", "num-integer", @@ -5691,9 +6101,9 @@ checksum = "0676bb32a98c1a483ce53e500a81ad9c3d5b3f7c920c28c24e9cb0980d0b5bc8" [[package]] name = "num" -version = "0.4.1" +version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b05180d69e3da0e530ba2a1dae5110317e49e3b7f3d41be227dc5f92e49ee7af" +checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" dependencies = [ "num-bigint", "num-complex", @@ -5705,11 +6115,10 @@ dependencies = [ [[package]] name = "num-bigint" -version = "0.4.4" +version = "0.4.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "608e7659b5c3d7cba262d894801b9ec9d00de989e8a82bd4bef91d08da45cdc0" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" dependencies = [ - "autocfg", "num-integer", "num-traits", ] @@ -5733,10 +6142,11 @@ dependencies = [ [[package]] name = "num-complex" -version = "0.4.5" +version = "0.4.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23c6602fda94a57c990fe0df199a035d83576b496aa29f4e634a8ac6004e68a6" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" dependencies = [ + "bytemuck", "num-traits", ] @@ -5759,19 +6169,18 @@ dependencies = [ [[package]] name = "num-integer" -version = "0.1.45" +version = "0.1.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" dependencies = [ - "autocfg", "num-traits", ] [[package]] name = "num-iter" -version = "0.1.43" +version = "0.1.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d03e6c028c5dc5cac6e2dec0efda81fc887605bb3d884578bb6d6bf7514e252" +checksum = "1429034a0490724d0075ebb2bc9e875d6503c3cf69e235a8941aa757d83ef5bf" dependencies = [ "autocfg", "num-integer", @@ -5780,11 +6189,10 @@ dependencies = [ [[package]] name = "num-rational" -version = "0.4.1" +version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0638a1c9d0a3c0914158145bc76cff373a75a627e6ecbfb71cbe6f453a5a19b0" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" dependencies = [ - "autocfg", "num-bigint", "num-integer", "num-traits", @@ -5841,6 +6249,12 @@ dependencies = [ "libc", ] +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + [[package]] name = "objc" version = "0.2.7" @@ -6083,6 +6497,28 @@ version = "1.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "945462a4b81e43c4e3ba96bd7b49d834c6f61198356aa858733bc4acf3cbe62e" +[[package]] +name = "onig" +version = "6.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c4b31c8722ad9171c6d77d3557db078cab2bd50afcc9d09c8b315c59df8ca4f" +dependencies = [ + "bitflags 1.3.2", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b829e3d7e9cc74c7e315ee8edb185bf4190da5acde74afd7fc59c35b1f086e7" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "opaque-debug" version = "0.3.0" @@ -6623,9 +7059,15 @@ dependencies = [ "crc32fast", "fdeflate", "flate2", - "miniz_oxide", + "miniz_oxide 0.7.2", ] +[[package]] +name = "portable-atomic" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "280dc24453071f1b63954171985a0b0d30058d287960968b9b2aca264c8d4ee6" + [[package]] name = "powerfmt" version = "0.2.0" @@ -6865,6 +7307,18 @@ version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "007d8adb5ddab6f8e3f491ac63566a7d5002cc7ed73901f72057943fa71ae1ae" +[[package]] +name = "pulp" +version = "0.18.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0a01a0dc67cf4558d279f0c25b0962bd08fc6dec0137699eae304103e882fe6" +dependencies = [ + "bytemuck", + "libm", + "num-complex", + "reborrow", +] + [[package]] name = "qoi" version = "0.4.1" @@ -7091,6 +7545,16 @@ dependencies = [ "getrandom 0.2.14", ] +[[package]] +name = "rand_distr" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31" +dependencies = [ + "num-traits", + "rand 0.8.5", +] + [[package]] name = "rand_hc" version = "0.2.0" @@ -7159,6 +7623,15 @@ dependencies = [ "rgb", ] +[[package]] +name = "raw-cpuid" +version = "10.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c297679cb867470fa8c9f67dbba74a78d78e3e98d7cf2b08d6d71540f797332" +dependencies = [ + "bitflags 1.3.2", +] + [[package]] name = "raw-window-handle" version = "0.6.2" @@ -7175,6 +7648,17 @@ dependencies = [ "rayon-core", ] +[[package]] +name = "rayon-cond" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "059f538b55efd2309c9794130bc149c6a553db90e9d99c2030785c82f0bd7df9" +dependencies = [ + "either", + "itertools 0.11.0", + "rayon", +] + [[package]] name = "rayon-core" version = "1.12.1" @@ -7194,6 +7678,12 @@ dependencies = [ "rand_core 0.3.1", ] +[[package]] +name = "reborrow" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03251193000f4bd3b042892be858ee50e8b3719f2b08e5833ac4353724632430" + [[package]] name = "redis" version = "0.13.0" @@ -7724,6 +8214,7 @@ version = "0.23.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2dabaac7466917e566adb06783a81ca48944c6898a1b08b9374106dd671f4c8" dependencies = [ + "log", "once_cell", "ring", "rustls-pki-types", @@ -7809,6 +8300,16 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" +[[package]] +name = "safetensors" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44560c11236a6130a46ce36c836a62936dc81ebf8c36a37947423571be0e55b6" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "same-file" version = "1.0.6" @@ -8013,6 +8514,12 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "388a1df253eca08550bef6c72392cfe7c30914bf41df5269b68cbd6ff8f570a3" +[[package]] +name = "seq-macro" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" + [[package]] name = "serde" version = "1.0.210" @@ -8517,6 +9024,18 @@ dependencies = [ "der", ] +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom", + "serde", + "unicode-segmentation", +] + [[package]] name = "sqlformat" version = "0.2.2" @@ -9213,6 +9732,31 @@ dependencies = [ "futures-core", ] +[[package]] +name = "synstructure" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.87", +] + +[[package]] +name = "sysctl" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec7dddc5f0fee506baf8b9fdb989e242f17e4b11c61dfbb0635b705217199eea" +dependencies = [ + "bitflags 2.5.0", + "byteorder", + "enum-as-inner", + "libc", + "thiserror 1.0.63", + "walkdir", +] + [[package]] name = "system-configuration" version = "0.5.1" @@ -9756,6 +10300,39 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tokenizers" +version = "0.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3dd47962b0ba36e7fd33518fbf1754d136fd1474000162bbf2a8b5fcb2d3654d" +dependencies = [ + "aho-corasick", + "clap", + "derive_builder", + "esaxx-rs", + "getrandom 0.2.14", + "indicatif", + "itertools 0.12.1", + "lazy_static", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand 0.8.5", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 1.0.63", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "tokens" version = "0.0.1" @@ -10250,6 +10827,33 @@ version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" +[[package]] +name = "ug" +version = "0.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4eef2ebfc18c67a6dbcacd9d8a4d85e0568cc58c82515552382312c2730ea13" +dependencies = [ + "half", + "num", + "serde", + "serde_json", + "thiserror 1.0.63", +] + +[[package]] +name = "ug-cuda" +version = "0.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c4dcab280ad0ef3957e153a82dcad608c954d02cf253b695322f502d1f8902e" +dependencies = [ + "cudarc", + "half", + "serde", + "serde_json", + "thiserror 1.0.63", + "ug", +] + [[package]] name = "unic-char-property" version = "0.9.0" @@ -10321,6 +10925,15 @@ dependencies = [ "tinyvec", ] +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec 1.13.2", +] + [[package]] name = "unicode-segmentation" version = "1.12.0" @@ -10333,6 +10946,12 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e51733f11c9c4f72aa0c160008246859e340b00807569a0da0e7a1079b27ba85" +[[package]] +name = "unicode-width" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd" + [[package]] name = "unicode_categories" version = "0.1.1" @@ -10354,6 +10973,25 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" +[[package]] +name = "ureq" +version = "2.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b74fc6b57825be3373f7054754755f03ac3a8f5d70015ccad699ba2029956f4a" +dependencies = [ + "base64 0.22.1", + "flate2", + "log", + "native-tls", + "once_cell", + "rustls 0.23.13", + "rustls-pki-types", + "serde", + "serde_json", + "url", + "webpki-roots 0.26.2", +] + [[package]] name = "url" version = "2.5.2" @@ -10775,6 +11413,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "web_scrapers" version = "0.0.1" @@ -11474,6 +12122,30 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "yoke" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "120e6aef9aa629e3d4f52dc8cc43a015c7724194c97dfaf45180d2daf2b77f40" +dependencies = [ + "serde", + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.87", + "synstructure", +] + [[package]] name = "zerocopy" version = "0.7.33" @@ -11494,6 +12166,27 @@ dependencies = [ "syn 2.0.87", ] +[[package]] +name = "zerofrom" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cff3ee08c995dee1859d998dea82f7374f2826091dd9cd47def953cae446cd2e" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "595eed982f7d355beb85837f651fa22e90b3c044842dc7f2c2842c086f295808" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.87", + "synstructure", +] + [[package]] name = "zeroize" version = "1.8.1" diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index 4e34dd0ef8..7dd9a5e191 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/crates/desktop/src-tauri/Cargo.toml @@ -26,3 +26,14 @@ tauri-plugin-log = "2.0.0-rc" once_cell = "1.20.3" image = "0.25.5" base64 = "0.22.1" +candle-core = { version = "0.8.2", features = ["cuda"] } +candle-transformers = { version = "0.8.2", features = ["cuda"] } +candle-nn = { version = "0.8.2", features = ["cuda"] } +tokenizers = "0.15.2" +hf-hub = "0.3.2" +anyhow = "1.0" +rand = "0.8" + +[features] +default = ["cuda"] +cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"] diff --git a/crates/desktop/src-tauri/src/image_endpoint.rs b/crates/desktop/src-tauri/src/image_endpoint.rs index a40e1a590f..028fc8e9cc 100644 --- a/crates/desktop/src-tauri/src/image_endpoint.rs +++ b/crates/desktop/src-tauri/src/image_endpoint.rs @@ -1,34 +1,51 @@ +use anyhow::Result; +use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; use std::io::Cursor; -use base64::Engine; -use base64::prelude::{BASE64_STANDARD, BASE64_URL_SAFE}; -use image::codecs::png::{CompressionType, FilterType, PngEncoder}; -use image::ImageReader; +use crate::sd::{self, StableDiffusionVersion, Args}; +use candle_core::Device; + +const PROMPT: &str = "A beautiful landscape with mountains and a lake"; #[tauri::command] pub fn infer_image(image: &str) -> Result { - println!("infer_image called; processing image..."); + println!("infer_image called; generating image with SDXL Turbo..."); + + // Check CUDA availability + let use_cpu = match Device::new_cuda(0) { + Ok(_) => { + println!("CUDA device available"); + false + } + Err(e) => { + println!("CUDA not available ({}), falling back to CPU", e); + true + } + }; + + let args = Args { + prompt: PROMPT.to_string(), + uncond_prompt: "".to_string(), + cpu: use_cpu, + height: Some(512), + width: Some(512), + n_steps: Some(1), + num_samples: 1, + seed: None, + sd_version: StableDiffusionVersion::Turbo, + guidance_scale: Some(0.0), + }; - let bytes = BASE64_STANDARD.decode(image) - .map_err(|err| format!("Base64 decode error: {}", err))?; - - let image = ImageReader::new(Cursor::new(bytes)) - .with_guessed_format() - .map_err(|err| format!("Image format error: {}", err))? - .decode() - .map_err(|err| format!("Image decode error: {}", err))?; - - let rotated_image = image.rotate180(); - - let mut buffer = Vec::new(); // TODO: Preallocate, better data structure - let mut writer = Cursor::new(&mut buffer); - - let encoder = PngEncoder::new_with_quality(&mut writer, CompressionType::Fast, FilterType::Adaptive); - rotated_image.write_with_encoder(encoder) - .map_err(|err| format!("Image encoding error: {}", err))?; - - let encoded = BASE64_STANDARD.encode(&buffer); - - println!("Encoded: {:?}", encoded.split_at(10).0); - - Ok(encoded) + match sd::run(args) { + Ok(_) => { + let img_data = std::fs::read("temp.png") + .map_err(|e| format!("Failed to read generated image: {}", e))?; + let encoded = BASE64_STANDARD.encode(&img_data); + println!("Generated image encoded successfully"); + + let _ = std::fs::remove_file("temp.png"); + + Ok(encoded) + } + Err(e) => Err(format!("Failed to generate image: {}", e)) + } } diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index eb94f6a947..4a9685fa39 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -1,4 +1,5 @@ mod image_endpoint; +mod sd; use std::sync::{Arc, RwLock}; use once_cell::sync::Lazy; diff --git a/crates/desktop/src-tauri/src/sd.rs b/crates/desktop/src-tauri/src/sd.rs new file mode 100644 index 0000000000..4aab7823e7 --- /dev/null +++ b/crates/desktop/src-tauri/src/sd.rs @@ -0,0 +1,385 @@ +#[cfg(feature = "accelerate")] +extern crate accelerate_src; + +#[cfg(feature = "mkl")] +extern crate intel_mkl_src; + +use candle_transformers::models::stable_diffusion; +use std::ops::Div; + +use anyhow::{Error as E, Result}; +use candle_core::{DType, Device, IndexOp, Tensor, D, Module}; +use image::{ImageBuffer, Rgb}; +use rand::Rng; +use stable_diffusion::vae::AutoEncoderKL; +use tokenizers::Tokenizer; + +// Helper function to save image +fn save_image>(img: &Tensor, p: P) -> Result<()> { + println!("Saving image to {:?}", p.as_ref()); + let p = p.as_ref(); + let (channel, height, width) = img.dims3()?; + println!("Image dimensions: {}x{} with {} channels", width, height, channel); + + if channel != 3 { + anyhow::bail!("save_image expects an input of shape (3, height, width)") + } + let img = img.permute((1, 2, 0))?.flatten_all()?; + let pixels = img.to_vec1::()?; + println!("Converting tensor to image buffer..."); + + let image: ImageBuffer, Vec> = + match ImageBuffer::from_raw(width as u32, height as u32, pixels) { + Some(image) => image, + None => anyhow::bail!("error saving image {p:?}"), + }; + println!("Successfully created image buffer"); + + image.save(p).map_err(E::from)?; + println!("Successfully saved image to disk"); + Ok(()) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StableDiffusionVersion { + V1_5, + V1_5Inpaint, + V2_1, + V2Inpaint, + Xl, + XlInpaint, + Turbo, +} + +pub struct Args { + pub prompt: String, + pub uncond_prompt: String, + pub cpu: bool, + pub height: Option, + pub width: Option, + pub n_steps: Option, + pub num_samples: usize, + pub seed: Option, + pub sd_version: StableDiffusionVersion, + pub guidance_scale: Option, +} + +impl StableDiffusionVersion { + fn repo(&self) -> &'static str { + match self { + Self::XlInpaint => "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", + Self::Xl => "stabilityai/stable-diffusion-xl-base-1.0", + Self::V2Inpaint => "stabilityai/stable-diffusion-2-inpainting", + Self::V2_1 => "stabilityai/stable-diffusion-2-1", + Self::V1_5 => "runwayml/stable-diffusion-v1-5", + Self::V1_5Inpaint => "stable-diffusion-v1-5/stable-diffusion-inpainting", + Self::Turbo => "stabilityai/sdxl-turbo", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ModelFile { + Tokenizer, + Tokenizer2, + Clip, + Clip2, + Unet, + Vae, +} + +impl ModelFile { + fn get( + &self, + filename: Option, + version: StableDiffusionVersion, + use_f16: bool, + ) -> Result { + use hf_hub::api::sync::Api; + match filename { + Some(filename) => Ok(std::path::PathBuf::from(filename)), + None => { + let (repo, path) = match self { + Self::Tokenizer => { + let tokenizer_repo = match version { + StableDiffusionVersion::V1_5 + | StableDiffusionVersion::V2_1 + | StableDiffusionVersion::V1_5Inpaint + | StableDiffusionVersion::V2Inpaint => "openai/clip-vit-base-patch32", + StableDiffusionVersion::Xl + | StableDiffusionVersion::XlInpaint + | StableDiffusionVersion::Turbo => { + "openai/clip-vit-large-patch14" + } + }; + (tokenizer_repo, "tokenizer.json") + } + Self::Tokenizer2 => { + ("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", "tokenizer.json") + } + Self::Clip => (version.repo(), "text_encoder/model.safetensors"), + Self::Clip2 => (version.repo(), "text_encoder_2/model.safetensors"), + Self::Unet => (version.repo(), "unet/diffusion_pytorch_model.safetensors"), + Self::Vae => (version.repo(), "vae/diffusion_pytorch_model.safetensors"), + }; + let filename = Api::new()?.model(repo.to_string()).get(path)?; + Ok(filename) + } + } + } +} + +pub fn run(args: Args) -> Result<()> { + println!("Starting image generation with the following configuration:"); + println!(" Model: {:?}", args.sd_version); + println!(" Prompt: {}", args.prompt); + println!(" Steps: {}", args.n_steps.unwrap_or(1)); + println!(" Using CPU: {}", args.cpu); + + let device = if args.cpu { + println!("Using CPU for computation"); + Device::Cpu + } else { + println!("Attempting to use CUDA device"); + match Device::new_cuda(0) { + Ok(cuda_device) => { + println!("Successfully initialized CUDA device"); + cuda_device + } + Err(e) => { + println!("Failed to initialize CUDA device: {}. Falling back to CPU", e); + Device::Cpu + } + } + }; + let dtype = DType::F32; + println!("Using {:?} data type", dtype); + + let sd_config = match args.sd_version { + StableDiffusionVersion::Turbo => { + println!("Configuring SDXL Turbo"); + stable_diffusion::StableDiffusionConfig::sdxl_turbo(None, args.height, args.width) + } + _ => { + println!("Configuring SD v2.1"); + stable_diffusion::StableDiffusionConfig::v2_1(None, args.height, args.width) + } + }; + println!("Model dimensions: {}x{}", sd_config.width, sd_config.height); + + println!("Building scheduler..."); + let mut scheduler = sd_config.build_scheduler(args.n_steps.unwrap_or(1))?; + let seed = args.seed.unwrap_or_else(|| rand::thread_rng().gen()); + println!("Using seed: {}", seed); + device.set_seed(seed)?; + + println!("Initializing Hugging Face API..."); + let api = hf_hub::api::sync::Api::new()?; + let repo = args.sd_version.repo(); + println!("Downloading model files from: {}", repo); + + println!("Downloading VAE and UNet..."); + let vae_file = api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?; + let unet_file = api.model(repo.to_string()).get("unet/diffusion_pytorch_model.safetensors")?; + + println!("Downloading text encoders..."); + let clip_file = api.model(repo.to_string()).get("text_encoder/model.safetensors")?; + let clip2_file = api.model(repo.to_string()).get("text_encoder_2/model.safetensors")?; + let tokenizer = api.model("openai/clip-vit-large-patch14".to_string()).get("tokenizer.json")?; + let tokenizer2 = api.model("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k".to_string()).get("tokenizer.json")?; + + println!("Building models..."); + let vae = sd_config.build_vae(vae_file, &device, dtype)?; + let unet = sd_config.build_unet(unet_file, &device, 4, false, dtype)?; + + // Build text encoders + println!("Building text encoders..."); + let text_model = stable_diffusion::build_clip_transformer(&sd_config.clip, clip_file, &device, dtype)?; + let text_model2 = stable_diffusion::build_clip_transformer(sd_config.clip2.as_ref().unwrap(), clip2_file, &device, dtype)?; + + // Tokenize and encode text + println!("Processing text prompt..."); + let tokenizer = tokenizers::Tokenizer::from_file(tokenizer).map_err(E::msg)?; + let tokenizer2 = tokenizers::Tokenizer::from_file(tokenizer2).map_err(E::msg)?; + + println!("Preparing text embeddings..."); + let which = match args.sd_version { + StableDiffusionVersion::Xl + | StableDiffusionVersion::XlInpaint + | StableDiffusionVersion::Turbo => vec![true, false], + _ => vec![true], + }; + + let text_embeddings = which + .iter() + .map(|first| { + text_embeddings( + &args.prompt, + &args.uncond_prompt, + None, // tokenizer + None, // clip_weights + None, // clip2_weights + args.sd_version, + &sd_config, + false, // use_f16 + &device, + dtype, + false, // use_guide_scale + *first, + ) + }) + .collect::>>()?; + + let text_embeddings = Tensor::cat(&text_embeddings, D::Minus1)?; + println!("Text embeddings shape: {:?}", text_embeddings.shape()); + + println!("Generating initial noise..."); + let timesteps = scheduler.timesteps().to_vec(); + let latents = Tensor::randn( + 0f32, + 1f32, + (1, 4, sd_config.height / 8, sd_config.width / 8), + &device, + )?; + let latents = (latents * scheduler.init_noise_sigma())?; + println!("Initial noise shape: {:?}", latents.shape()); + + println!("Starting diffusion process..."); + let mut latents = latents; + for (timestep_index, ×tep) in timesteps.iter().enumerate() { + println!("Processing step {}/{}", timestep_index + 1, timesteps.len()); + let latent_model_input = latents.clone(); + println!("Latent input shape: {:?}", latent_model_input.shape()); + + println!("Scaling model input..."); + let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?; + println!("Scaled input shape: {:?}", latent_model_input.shape()); + + println!("Running UNet inference with timestep {}...", timestep); + let noise_pred = match unet.forward(&latent_model_input, timestep as f64, &text_embeddings) { + Ok(pred) => { + println!("UNet inference successful"); + pred + } + Err(e) => { + println!("UNet inference failed with error: {}", e); + return Err(anyhow::anyhow!("UNet inference failed: {}", e)); + } + }; + println!("Noise prediction shape: {:?}", noise_pred.shape()); + + println!("Applying scheduler step..."); + latents = scheduler.step(&noise_pred, timestep, &latents)?; + println!("Step {}/{} completed", timestep_index + 1, timesteps.len()); + } + + println!("Diffusion process completed, decoding image..."); + let image = vae.decode(&(latents / 0.13025)?)?; + println!("VAE decode completed"); + let image = ((image / 2.)? + 0.5)?.to_device(&Device::Cpu)?; + println!("Normalized image values"); + let image = (image.clamp(0f32, 1.)? * 255.)?.to_dtype(DType::U8)?; + println!("Converted to 8-bit format"); + + save_image(&image.i(0)?, "temp.png")?; + println!("Image generation completed successfully"); + + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +fn text_embeddings( + prompt: &str, + uncond_prompt: &str, + tokenizer: Option, + clip_weights: Option, + clip2_weights: Option, + sd_version: StableDiffusionVersion, + sd_config: &stable_diffusion::StableDiffusionConfig, + use_f16: bool, + device: &Device, + dtype: DType, + use_guide_scale: bool, + first: bool, +) -> Result { + let tokenizer_file = if first { + ModelFile::Tokenizer + } else { + ModelFile::Tokenizer2 + }; + let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?; + let tokenizer = Tokenizer::from_file(tokenizer).map_err(E::msg)?; + let pad_id = match &sd_config.clip.pad_with { + Some(padding) => *tokenizer.get_vocab(true).get(padding.as_str()).unwrap(), + None => *tokenizer.get_vocab(true).get("<|endoftext|>").unwrap(), + }; + + println!("Running with prompt \"{prompt}\"."); + let mut tokens = tokenizer + .encode(prompt, true) + .map_err(E::msg)? + .get_ids() + .to_vec(); + if tokens.len() > sd_config.clip.max_position_embeddings { + anyhow::bail!( + "the prompt is too long, {} > max-tokens ({})", + tokens.len(), + sd_config.clip.max_position_embeddings + ) + } + while tokens.len() < sd_config.clip.max_position_embeddings { + tokens.push(pad_id) + } + let tokens = Tensor::new(tokens.as_slice(), device)?.unsqueeze(0)?; + + println!("Building the Clip transformer."); + let clip_weights_file = if first { + ModelFile::Clip + } else { + ModelFile::Clip2 + }; + let clip_weights = if first { + clip_weights_file.get(clip_weights, sd_version, use_f16)? + } else { + clip_weights_file.get(clip2_weights, sd_version, use_f16)? + }; + let clip_config = if first { + &sd_config.clip + } else { + sd_config.clip2.as_ref().unwrap() + }; + let text_model = stable_diffusion::build_clip_transformer(clip_config, clip_weights, device, DType::F32)?; + + let text_embeddings = text_model.forward(&tokens)?; + + let text_embeddings = if use_guide_scale { + let mut uncond_tokens = tokenizer + .encode(uncond_prompt, true) + .map_err(E::msg)? + .get_ids() + .to_vec(); + if uncond_tokens.len() > sd_config.clip.max_position_embeddings { + anyhow::bail!( + "the negative prompt is too long, {} > max-tokens ({})", + uncond_tokens.len(), + sd_config.clip.max_position_embeddings + ) + } + while uncond_tokens.len() < sd_config.clip.max_position_embeddings { + uncond_tokens.push(pad_id) + } + + let uncond_tokens = Tensor::new(uncond_tokens.as_slice(), device)?.unsqueeze(0)?; + let uncond_embeddings = text_model.forward(&uncond_tokens)?; + + Tensor::cat(&[uncond_embeddings, text_embeddings], 0)?.to_dtype(dtype)? + } else { + text_embeddings.to_dtype(dtype)? + }; + Ok(text_embeddings) +} + +// fn main() -> Result<()> { +// let args = Args::parse(); +// run(args) +// }