From e3944588fc1d88ccfc58248022f8be45762a1078 Mon Sep 17 00:00:00 2001 From: jhqxxx <18280426169@163.com> Date: Fri, 30 Jan 2026 22:04:23 +0800 Subject: [PATCH] index tts stash save --- Cargo.lock | 742 +++++++++++++++++- Cargo.toml | 5 +- README.md | 2 +- download_and_run.sh | 4 +- src/main.rs | 62 +- src/models/campplus/mod.rs | 549 +++++++++++++ src/models/common/mod.rs | 341 +++++++- src/models/deepseek_ocr/model.rs | 9 + src/models/feature_extractor/mod.rs | 1 + .../seamless_m4t_feature_extractor.rs | 129 +++ src/models/fun_asr_nano/model.rs | 6 +- src/models/glm_asr_nano/processor.rs | 23 +- src/models/hunyuan_ocr/model.rs | 3 + src/models/index_tts2/config.rs | 228 ++++++ src/models/index_tts2/generate.rs | 29 + src/models/index_tts2/mod.rs | 5 + src/models/index_tts2/model.rs | 617 +++++++++++++++ src/models/index_tts2/processor.rs | 199 +++++ src/models/index_tts2/utils.rs | 18 + src/models/mask_gct/config.rs | 21 + src/models/mask_gct/mod.rs | 2 + src/models/mask_gct/model.rs | 344 ++++++++ src/models/minicpm4/model.rs | 3 + src/models/mod.rs | 7 +- src/models/qwen2_5vl/model.rs | 6 + src/models/qwen3/model.rs | 3 + src/models/rmbg2_0/generate.rs | 2 +- src/models/rmbg2_0/model.rs | 49 +- src/models/voxcpm/minicpm4.rs | 3 + src/models/w2v_bert_2_0/config.rs | 61 ++ src/models/w2v_bert_2_0/mod.rs | 2 + src/models/w2v_bert_2_0/model.rs | 571 ++++++++++++++ src/position_embed/rope.rs | 64 +- src/utils/audio_utils.rs | 125 ++- src/utils/mod.rs | 274 ++++++- src/utils/tensor_utils.rs | 136 +++- tests/messy_test.rs | 17 +- tests/test_index_tts2.rs | 48 ++ tests/weight_test.rs | 40 +- 39 files changed, 4613 insertions(+), 137 deletions(-) create mode 100644 src/models/campplus/mod.rs create mode 100644 src/models/feature_extractor/mod.rs create mode 100644 src/models/feature_extractor/seamless_m4t_feature_extractor.rs create mode 100644 src/models/index_tts2/config.rs create mode 100644 src/models/index_tts2/generate.rs create mode 100644 src/models/index_tts2/mod.rs create mode 100644 src/models/index_tts2/model.rs create mode 100644 src/models/index_tts2/processor.rs create mode 100644 src/models/index_tts2/utils.rs create mode 100644 src/models/mask_gct/config.rs create mode 100644 src/models/mask_gct/mod.rs create mode 100644 src/models/mask_gct/model.rs create mode 100644 src/models/w2v_bert_2_0/config.rs create mode 100644 src/models/w2v_bert_2_0/mod.rs create mode 100644 src/models/w2v_bert_2_0/model.rs create mode 100644 tests/test_index_tts2.rs diff --git a/Cargo.lock b/Cargo.lock index d09013d..1791314 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,17 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + [[package]] name = "aha" version = "0.1.8" @@ -24,6 +35,7 @@ dependencies = [ "aha_openai_dive", "anyhow", "base64 0.22.1", + "byteorder", "candle-core", "candle-flash-attn", "candle-nn", @@ -32,6 +44,7 @@ dependencies = [ "clap", "dirs", "ffmpeg-next", + "half", "hound", "image", "minijinja", @@ -39,7 +52,7 @@ dependencies = [ "num", "rayon", "realfft", - "reqwest", + "reqwest 0.12.24", "rocket", "serde", "serde_json", @@ -49,6 +62,7 @@ dependencies = [ "tokio", "url", "uuid", + "zip 7.2.0", ] [[package]] @@ -60,7 +74,7 @@ dependencies = [ "bytes", "derive_builder", "futures", - "reqwest", + "reqwest 0.12.24", "reqwest-eventsource", "serde", "serde_html_form", @@ -279,6 +293,28 @@ dependencies = [ "arrayvec", ] +[[package]] +name = "aws-lc-rs" +version = "1.15.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e84ce723ab67259cfeb9877c6a639ee9eb7a27b28123abd71db7f0d5d0cc9d86" +dependencies = [ + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.36.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43a442ece363113bd4bd4c8b18977a7798dd4d3c3383f34fb61936960e8f4ad8" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", +] + [[package]] name = "backtrace" version = "0.3.76" @@ -386,6 +422,15 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "built" version = "0.7.7" @@ -439,6 +484,15 @@ dependencies = [ "serde", ] +[[package]] +name = "bzip2" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3a53fac24f34a81bc9954b5d6cfce0c21e18ec6959f44f56e8e90e4bb7c346c" +dependencies = [ + "libbz2-rs-sys", +] + [[package]] name = "candle-core" version = "0.9.1" @@ -464,7 +518,7 @@ dependencies = [ "ug-cuda", "ug-metal", "yoke 0.7.5", - "zip", + "zip 1.1.4", ] [[package]] @@ -548,9 +602,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.43" +version = "1.2.53" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "739eb0f94557554b3ca9a86d2d37bebd49c5e6d0c1d2bda35ba5bdac830befc2" +checksum = "755d2fce177175ffca841e9a06afdb2c4ab0f593d53b4dee48147dfaade85932" dependencies = [ "find-msvc-tools", "jobserver", @@ -558,6 +612,12 @@ dependencies = [ "shlex", ] +[[package]] +name = "cesu8" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c" + [[package]] name = "cexpr" version = "0.6.0" @@ -583,6 +643,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "chrono" version = "0.4.42" @@ -596,6 +662,16 @@ dependencies = [ "windows-link 0.2.1", ] +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common", + "inout", +] + [[package]] name = "clang-sys" version = "1.8.1" @@ -647,6 +723,15 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a1d728cc89cf3aee9ff92b05e62b19ee65a02b5702cff7d5a377e32c6ae29d8d" +[[package]] +name = "cmake" +version = "0.1.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75443c44cd6b379beb8c5b45d85d0773baf31cce901fe7bb252f4eff3008ef7d" +dependencies = [ + "cc", +] + [[package]] name = "color_quant" version = "1.1.0" @@ -659,6 +744,16 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" +[[package]] +name = "combine" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +dependencies = [ + "bytes", + "memchr", +] + [[package]] name = "compact_str" version = "0.9.0" @@ -700,6 +795,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "constant_time_eq" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c74b8349d32d297c9134b8c88677813a227df8f779daa29bfc29c183fe3dca6" + [[package]] name = "cookie" version = "0.18.1" @@ -711,6 +812,24 @@ dependencies = [ "version_check", ] +[[package]] +name = "cookie_store" +version = "0.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fc4bff745c9b4c7fb1e97b25d13153da2bc7796260141df62378998d070207f" +dependencies = [ + "cookie", + "document-features", + "idna", + "log", + "publicsuffix", + "serde", + "serde_derive", + "serde_json", + "time", + "url", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -721,6 +840,16 @@ dependencies = [ "libc", ] +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "core-foundation-sys" version = "0.8.7" @@ -734,10 +863,34 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "45390e6114f68f718cc7a830514a96f903cccd70d02a8f6d9f643ac4ba45afaf" dependencies = [ "bitflags 1.3.2", - "core-foundation", + "core-foundation 0.9.4", "libc", ] +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc" +version = "3.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9710d3b3739c2e349eb44fe848ad0b7c8cb1e42bd87ee49371df2f7acaf3e675" +dependencies = [ + "crc-catalog", +] + +[[package]] +name = "crc-catalog" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19d374276b40fb8bbdee95aef7c7fa6b5316ec764510eb64b8dd0e2ed0d7e7f5" + [[package]] name = "crc32fast" version = "1.5.0" @@ -778,6 +931,16 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + [[package]] name = "cudarc" version = "0.16.6" @@ -832,6 +995,12 @@ dependencies = [ "serde", ] +[[package]] +name = "deflate64" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26bf8fc351c5ed29b5c2f0cbbac1b209b74f60ecd62e675a998df72c49af5204" + [[package]] name = "deranged" version = "0.5.5" @@ -916,6 +1085,17 @@ dependencies = [ "syn", ] +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", + "subtle", +] + [[package]] name = "dirs" version = "6.0.0" @@ -948,6 +1128,21 @@ dependencies = [ "syn", ] +[[package]] +name = "document-features" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61" +dependencies = [ + "litrs", +] + +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + [[package]] name = "dyn-stack" version = "0.10.0" @@ -1171,9 +1366,9 @@ dependencies = [ [[package]] name = "find-msvc-tools" -version = "0.1.4" +version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52051878f80a721bb68ebfbc930e07b65ba72f2da88968ea5c06fd6ca3d3a127" +checksum = "8591b0bcc8a98a64310a2fae1bb3e9b8564dd10e381e6e28010fde8e8e8568db" [[package]] name = "flate2" @@ -1182,6 +1377,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bfe33edd8e85a12a67454e37f8c75e730830d83e313556ab9ebf9ee7fbeb3bfb" dependencies = [ "crc32fast", + "libz-rs-sys", "miniz_oxide", ] @@ -1242,6 +1438,12 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + [[package]] name = "futures" version = "0.3.31" @@ -1587,6 +1789,16 @@ dependencies = [ "windows", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "getrandom" version = "0.2.16" @@ -1594,8 +1806,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" dependencies = [ "cfg-if", + "js-sys", "libc", "wasi", + "wasm-bindgen", ] [[package]] @@ -1605,9 +1819,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi", "wasip2", + "wasm-bindgen", ] [[package]] @@ -1703,6 +1919,15 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest", +] + [[package]] name = "hound" version = "3.5.1" @@ -2102,6 +2327,15 @@ version = "0.1.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8fae54786f62fb2918dcfae3d568594e50eb9b5c25bf04371af6fe7516452fb" +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + [[package]] name = "interpolate_name" version = "0.2.4" @@ -2179,6 +2413,28 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" +[[package]] +name = "jni" +version = "0.21.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a87aa2bb7d2af34197c04845522473242e1aa17c12f4935d5856491a7fb8c97" +dependencies = [ + "cesu8", + "cfg-if", + "combine", + "jni-sys", + "log", + "thiserror 1.0.69", + "walkdir", + "windows-sys 0.45.0", +] + +[[package]] +name = "jni-sys" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8eaf4bc02d17cbdd7ff4c7438cafcdf7fb9a4613313ad11b4f8fefe7d3fa0130" + [[package]] name = "jobserver" version = "0.1.34" @@ -2211,6 +2467,12 @@ version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7a79a3332a6609480d7d0c9eab957bca6b455b91bb84e66d19f5ff66294b85b8" +[[package]] +name = "libbz2-rs-sys" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c4a545a15244c7d945065b5d392b2d2d7f21526fba56ce51467b06ed445e8f7" + [[package]] name = "libc" version = "0.2.177" @@ -2253,6 +2515,15 @@ dependencies = [ "libc", ] +[[package]] +name = "libz-rs-sys" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c10501e7805cee23da17c7790e59df2870c0d4043ec6d03f67d31e2b53e77415" +dependencies = [ + "zlib-rs", +] + [[package]] name = "linux-raw-sys" version = "0.11.0" @@ -2265,6 +2536,12 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "241eaef5fd12c88705a01fc1066c48c4b36e0dd4377dcdc7ec3942cea7a69956" +[[package]] +name = "litrs" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" + [[package]] name = "lock_api" version = "0.4.14" @@ -2304,6 +2581,22 @@ dependencies = [ "imgref", ] +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + +[[package]] +name = "lzma-rust2" +version = "0.15.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1670343e58806300d87950e3401e820b519b9384281bbabfb15e3636689ffd69" +dependencies = [ + "crc", + "sha2", +] + [[package]] name = "macro_rules_attribute" version = "0.2.2" @@ -2448,17 +2741,17 @@ dependencies = [ [[package]] name = "modelscope" -version = "0.1.0" +version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe4bcb2259d8c3b94cb7041a55dcca7a5d16de888b45f7590e389f746e52e560" +checksum = "43f5192d2b709f323d867523bde2bd7b9536ba701f7eb7e3c938a74193e2791d" dependencies = [ "anyhow", "clap", "futures-util", "indicatif 0.18.2", - "openssl", - "reqwest", + "reqwest 0.13.1", "serde", + "serde_json", "tokio", ] @@ -2522,10 +2815,10 @@ dependencies = [ "libc", "log", "openssl", - "openssl-probe", + "openssl-probe 0.1.6", "openssl-sys", "schannel", - "security-framework", + "security-framework 2.11.1", "security-framework-sys", "tempfile", ] @@ -2795,13 +3088,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e" [[package]] -name = "openssl-src" -version = "300.5.4+3.5.4" +name = "openssl-probe" +version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a507b3792995dae9b0df8a1c1e3771e8418b7c2d9f0baeba32e6fe8b06c7cb72" -dependencies = [ - "cc", -] +checksum = "9f50d9b3dabb09ecd771ad0aa242ca6894994c130308ca3d7684634df8037391" [[package]] name = "openssl-sys" @@ -2811,7 +3101,6 @@ checksum = "0a9f0075ba3c21b09f8e8b2026584b1d18d49388648f2fbbf3c97ea8deced8e2" dependencies = [ "cc", "libc", - "openssl-src", "pkg-config", "vcpkg", ] @@ -2851,6 +3140,16 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pbkdf2" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2" +dependencies = [ + "digest", + "hmac", +] + [[package]] name = "pear" version = "0.2.9" @@ -2932,6 +3231,12 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" +[[package]] +name = "ppmd-rust" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "efca4c95a19a79d1c98f791f10aebd5c1363b473244630bb7dbde1dc98455a24" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -3000,6 +3305,22 @@ dependencies = [ "syn", ] +[[package]] +name = "psl-types" +version = "2.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33cb294fe86a74cbcf50d4445b37da762029549ebeea341421c7c70370f86cac" + +[[package]] +name = "publicsuffix" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f42ea446cab60335f76979ec15e12619a2165b5ae2c12166bef27d283a9fadf" +dependencies = [ + "idna", + "psl-types", +] + [[package]] name = "pulp" version = "0.18.22" @@ -3050,6 +3371,62 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a993555f31e5a609f617c12db6250dedcac1b0a85076912c436e6fc9b2c8e6a3" +[[package]] +name = "quinn" +version = "0.11.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash", + "rustls", + "socket2 0.6.1", + "thiserror 2.0.17", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1906b49b0c3bc04b5fe5d86a77925ae6524a19b816ae38ce1e426255f1d8a31" +dependencies = [ + "aws-lc-rs", + "bytes", + "getrandom 0.3.4", + "lru-slab", + "rand 0.9.2", + "ring", + "rustc-hash", + "rustls", + "rustls-pki-types", + "slab", + "thiserror 2.0.17", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2 0.6.1", + "tracing", + "windows-sys 0.60.2", +] + [[package]] name = "quote" version = "1.0.41" @@ -3362,6 +3739,51 @@ dependencies = [ "web-sys", ] +[[package]] +name = "reqwest" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04e9018c9d814e5f30cc16a0f03271aeab3571e609612d9fe78c1aa8d11c2f62" +dependencies = [ + "base64 0.22.1", + "bytes", + "cookie", + "cookie_store", + "encoding_rs", + "futures-core", + "futures-util", + "h2 0.4.12", + "http 1.3.1", + "http-body 1.0.1", + "http-body-util", + "hyper 1.7.0", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "mime", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-pki-types", + "rustls-platform-verifier", + "serde", + "serde_json", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", +] + [[package]] name = "reqwest-eventsource" version = "0.6.0" @@ -3374,7 +3796,7 @@ dependencies = [ "mime", "nom 7.1.3", "pin-project-lite", - "reqwest", + "reqwest 0.12.24", "thiserror 1.0.69", ] @@ -3525,6 +3947,7 @@ version = "0.23.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a9586e9ee2b4f8fab52a0048ca7334d7024eef48e2cb9407e3497bb7cab7fa7" dependencies = [ + "aws-lc-rs", "once_cell", "rustls-pki-types", "rustls-webpki", @@ -3532,21 +3955,62 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-native-certs" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63" +dependencies = [ + "openssl-probe 0.2.0", + "rustls-pki-types", + "schannel", + "security-framework 3.5.1", +] + [[package]] name = "rustls-pki-types" version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "229a4a4c221013e7e1f1a043678c5cc39fe5171437c88fb47151a21e6f5b5c79" dependencies = [ + "web-time", "zeroize", ] +[[package]] +name = "rustls-platform-verifier" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d99feebc72bae7ab76ba994bb5e121b8d83d910ca40b36e0921f53becc41784" +dependencies = [ + "core-foundation 0.10.1", + "core-foundation-sys", + "jni", + "log", + "once_cell", + "rustls", + "rustls-native-certs", + "rustls-platform-verifier-android", + "rustls-webpki", + "security-framework 3.5.1", + "security-framework-sys", + "webpki-root-certs", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls-platform-verifier-android" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" + [[package]] name = "rustls-webpki" version = "0.103.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e10b3f4191e8a80e6b43eebabfac91e5dcecebb27a71f04e820c47ec41d314bf" dependencies = [ + "aws-lc-rs", "ring", "rustls-pki-types", "untrusted", @@ -3611,7 +4075,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "897b2245f0b511c87893af39b033e5ca9cce68824c4d7e7630b5a1d339658d02" dependencies = [ "bitflags 2.10.0", - "core-foundation", + "core-foundation 0.9.4", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework" +version = "3.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3297343eaf830f66ede390ea39da1d462b6b0c1b000f420d0a83f898bbbe6ef" +dependencies = [ + "bitflags 2.10.0", + "core-foundation 0.10.1", "core-foundation-sys", "libc", "security-framework-sys", @@ -3732,6 +4209,28 @@ dependencies = [ "unsafe-libyaml", ] +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "sharded-slab" version = "0.1.7" @@ -4081,7 +4580,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c879d448e9d986b661742763247d3693ed13609438cf3d006f51f5368a5ba6b" dependencies = [ "bitflags 2.10.0", - "core-foundation", + "core-foundation 0.9.4", "system-configuration-sys", ] @@ -4198,6 +4697,7 @@ checksum = "91e7d9e3bb61134e77bde20dd4825b97c010155709965fedf0f49bb138e52a9d" dependencies = [ "deranged", "itoa", + "js-sys", "num-conv", "powerfmt", "serde", @@ -4231,6 +4731,21 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinyvec" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa5fdc3bce6191a1dbc8c02d5c8bffcf557bafa17c124c5264a458f1b0613fa" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + [[package]] name = "tokenizers" version = "0.22.1" @@ -4424,9 +4939,9 @@ dependencies = [ [[package]] name = "tower-http" -version = "0.6.6" +version = "0.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adc82fd73de2a9722ac5da747f12383d2bfdb93591ee6c58486e0097890f05f2" +checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" dependencies = [ "bitflags 2.10.0", "bytes", @@ -4529,6 +5044,18 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "typed-path" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3015e6ce46d5ad8751e4a772543a30c7511468070e98e64e20165f8f81155b64" + +[[package]] +name = "typenum" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb" + [[package]] name = "ubyte" version = "0.10.4" @@ -4868,6 +5395,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-root-certs" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36a29fc0408b113f68cf32637857ab740edfafdf460c326cd2afaa2d84cc05dc" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "weezl" version = "0.1.10" @@ -4986,6 +5522,15 @@ dependencies = [ "windows-link 0.2.1", ] +[[package]] +name = "windows-sys" +version = "0.45.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0" +dependencies = [ + "windows-targets 0.42.2", +] + [[package]] name = "windows-sys" version = "0.52.0" @@ -5022,6 +5567,21 @@ dependencies = [ "windows-link 0.2.1", ] +[[package]] +name = "windows-targets" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e5180c00cd44c9b1c88adb3693291f1cd93605ded80c250a75d472756b4d071" +dependencies = [ + "windows_aarch64_gnullvm 0.42.2", + "windows_aarch64_msvc 0.42.2", + "windows_i686_gnu 0.42.2", + "windows_i686_msvc 0.42.2", + "windows_x86_64_gnu 0.42.2", + "windows_x86_64_gnullvm 0.42.2", + "windows_x86_64_msvc 0.42.2", +] + [[package]] name = "windows-targets" version = "0.48.5" @@ -5070,6 +5630,12 @@ dependencies = [ "windows_x86_64_msvc 0.53.1", ] +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" + [[package]] name = "windows_aarch64_gnullvm" version = "0.48.5" @@ -5088,6 +5654,12 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" +[[package]] +name = "windows_aarch64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" + [[package]] name = "windows_aarch64_msvc" version = "0.48.5" @@ -5106,6 +5678,12 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" +[[package]] +name = "windows_i686_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" + [[package]] name = "windows_i686_gnu" version = "0.48.5" @@ -5136,6 +5714,12 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" +[[package]] +name = "windows_i686_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" + [[package]] name = "windows_i686_msvc" version = "0.48.5" @@ -5154,6 +5738,12 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" +[[package]] +name = "windows_x86_64_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" + [[package]] name = "windows_x86_64_gnu" version = "0.48.5" @@ -5172,6 +5762,12 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" + [[package]] name = "windows_x86_64_gnullvm" version = "0.48.5" @@ -5190,6 +5786,12 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" +[[package]] +name = "windows_x86_64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" + [[package]] name = "windows_x86_64_msvc" version = "0.48.5" @@ -5332,6 +5934,20 @@ name = "zeroize" version = "1.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85a5b4158499876c763cb03bc4e49185d3cccbabb15b33c627f7884f43db852e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] [[package]] name = "zerotrie" @@ -5381,6 +5997,80 @@ dependencies = [ "thiserror 1.0.69", ] +[[package]] +name = "zip" +version = "7.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c42e33efc22a0650c311c2ef19115ce232583abbe80850bc8b66509ebef02de0" +dependencies = [ + "aes", + "bzip2", + "constant_time_eq", + "crc32fast", + "deflate64", + "flate2", + "generic-array", + "getrandom 0.3.4", + "hmac", + "indexmap", + "lzma-rust2", + "memchr", + "pbkdf2", + "ppmd-rust", + "sha1", + "time", + "typed-path", + "zeroize", + "zopfli", + "zstd", +] + +[[package]] +name = "zlib-rs" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40990edd51aae2c2b6907af74ffb635029d5788228222c4bb811e9351c0caad3" + +[[package]] +name = "zopfli" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f05cd8797d63865425ff89b5c4a48804f35ba0ce8d125800027ad6017d2b5249" +dependencies = [ + "bumpalo", + "crc32fast", + "log", + "simd-adler32", +] + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "zune-core" version = "0.4.12" diff --git a/Cargo.toml b/Cargo.toml index 2f5bdbd..3f5c6e3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,7 +28,7 @@ rocket = { version = "0.5.1", features = ["serde_json", "json"] } tokio = "1.47.1" hound = "3.5.1" clap = { version = "4.5.51", features = ["derive"] } -modelscope = "0.1.0" +modelscope = "0.1.3" dirs = "6.0.0" url = "2.5.7" rayon = "1.10" @@ -37,6 +37,9 @@ rayon = "1.10" realfft = "3.5.0" symphonia = { version = "0.5.5", features = ["mp3", "wav"] } serde_yaml = "0.9.34" +zip = "7.2.0" +half = "2.7.1" +byteorder = "1.5.0" [features] flash-attn = ["candle-flash-attn"] diff --git a/README.md b/README.md index d3a959d..12aa4bf 100644 --- a/README.md +++ b/README.md @@ -118,7 +118,7 @@ cargo run -F cuda -r -- [参数] * deepseek-ocr: deepseek-ai/DeepSeek-OCR 模型 * hunyuan-ocr: Tencent-Hunyuan/HunyuanOCR 模型 * paddleocr-vl: PaddlePaddle/PaddleOCR-VL 模型 - * RMBG2.0: AI-ModelScope/RMBG-2.0 模型 + * rmbg2.0: AI-ModelScope/RMBG-2.0 模型 * voxcpm: OpenBMB/VoxCPM-0.5B 模型 * voxcpm1.5: OpenBMB/VoxCPM1.5 模型 * glm-asr-nano-2512: ZhipuAI/GLM-ASR-Nano-2512 模型 diff --git a/download_and_run.sh b/download_and_run.sh index 684aab1..814667f 100755 --- a/download_and_run.sh +++ b/download_and_run.sh @@ -25,7 +25,7 @@ show_help() { echo " deepseek-ocr" echo " hunyuan-ocr" echo " paddleocr-vl" - echo " RMBG2.0" + echo " rmbg2.0" echo " voxcpm" echo " voxcpm1.5" echo " glm-asr-nano-2512" @@ -77,7 +77,7 @@ case $MODEL_ALIAS in "paddleocr-vl") MODEL_ID="PaddlePaddle/PaddleOCR-VL" ;; - "RMBG2.0") + "rmbg2.0") MODEL_ID="briaai/RMBG-2.0" ;; "voxcpm") diff --git a/src/main.rs b/src/main.rs index be1118e..0a1cb22 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,6 @@ use std::{net::IpAddr, str::FromStr, time::Duration}; -use aha::{models::WhichModel, utils::get_default_save_dir}; +use aha::{models::WhichModel, utils::{download_model, get_default_save_dir}}; use clap::Parser; use modelscope::ModelScope; use rocket::{ @@ -34,38 +34,38 @@ struct Args { #[arg(long)] download_retries: Option, } -async fn download_model(model_id: &str, save_dir: &str, max_retries: u32) -> anyhow::Result<()> { - let mut attempts = 0u32; - loop { - attempts += 1; - println!( - "Attempting to download model (attempt {}/{})", - attempts, max_retries - ); +// async fn download_model(model_id: &str, save_dir: &str, max_retries: u32) -> anyhow::Result<()> { +// let mut attempts = 0u32; +// loop { +// attempts += 1; +// println!( +// "Attempting to download model (attempt {}/{})", +// attempts, max_retries +// ); - match ModelScope::download(model_id, save_dir).await { - Ok(()) => { - println!("Model downloaded successfully"); - return Ok(()); - } - Err(e) => { - if attempts >= max_retries { - return Err(anyhow::anyhow!( - "Failed to download model after {} attempts. Last error: {}", - max_retries, - e - )); - } +// match ModelScope::download(model_id, save_dir).await { +// Ok(()) => { +// println!("Model downloaded successfully"); +// return Ok(()); +// } +// Err(e) => { +// if attempts >= max_retries { +// return Err(anyhow::anyhow!( +// "Failed to download model after {} attempts. Last error: {}", +// max_retries, +// e +// )); +// } - println!( - "Download failed (attempt {}): {}. Retrying in 2 seconds...", - attempts, e - ); - sleep(Duration::from_secs(2)).await; - } - } - } -} +// println!( +// "Download failed (attempt {}): {}. Retrying in 2 seconds...", +// attempts, e +// ); +// sleep(Duration::from_secs(2)).await; +// } +// } +// } +// } #[tokio::main] async fn main() -> anyhow::Result<()> { diff --git a/src/models/campplus/mod.rs b/src/models/campplus/mod.rs new file mode 100644 index 0000000..d5638e5 --- /dev/null +++ b/src/models/campplus/mod.rs @@ -0,0 +1,549 @@ +use anyhow::Result; +use candle_core::{D, Tensor}; +use candle_nn::{BatchNorm, Conv1d, Conv2d, Module, ModuleT, VarBuilder, ops::sigmoid}; + +use crate::{ + models::common::{get_batch_norm, get_conv1d, get_conv2d}, + utils::tensor_utils::{pool1d, statistics_pooling}, +}; + +pub struct Shortcut { + conv_0: Conv2d, + bn_1: BatchNorm, + stride: usize, +} + +impl Shortcut { + pub fn new( + vb: VarBuilder, + in_c: usize, + out_c: usize, + ks: usize, + padding: usize, + stride: usize, + bias: bool, + ) -> Result { + let conv_0 = get_conv2d(vb.pp("0"), in_c, out_c, ks, padding, 1, 1, 1, bias)?; + let bn_1 = get_batch_norm(vb.pp("1"), 1e-5, out_c, true)?; + Ok(Self { conv_0, bn_1, stride }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let mut x = self.conv_0.forward(x)?; + if self.stride != 1 { + let h_dim = x.dim(2)?; + let half_h = h_dim / 2; + let indices = Tensor::arange(0u32, half_h as u32, x.device())?.affine(2.0, 0.0)?; + x = x.index_select(&indices, 2)?; + } + x = self.bn_1.forward_t(&x, false)?; + Ok(x) + } +} + +pub struct BasicResBlock { + stride: usize, + conv1: Conv2d, + bn1: BatchNorm, + conv2: Conv2d, + bn2: BatchNorm, + shortcut: Option, +} + +impl BasicResBlock { + pub fn new( + vb: VarBuilder, + in_planes: usize, + planes: usize, + stride: usize, + expansion: usize, + ) -> Result { + let conv1 = get_conv2d(vb.pp("conv1"), in_planes, planes, 3, 1, 1, 1, 1, false)?; + let bn1 = get_batch_norm(vb.pp("bn1"), 1e-5, planes, true)?; + let conv2 = get_conv2d(vb.pp("conv2"), planes, planes, 3, 1, 1, 1, 1, false)?; + let bn2 = get_batch_norm(vb.pp("bn2"), 1e-5, planes, true)?; + let shortcut = if stride != 1 || in_planes != expansion * planes { + Some(Shortcut::new( + vb.pp("shortcut"), + in_planes, + expansion * planes, + 1, + 0, + stride, + false, + )?) + } else { + None + }; + Ok(Self { + stride, + conv1, + bn1, + conv2, + bn2, + shortcut, + }) + } + pub fn forward(&self, xs: &Tensor) -> Result { + let residual = xs.clone(); + let mut xs = self.conv1.forward(xs)?; + // candle stride only surpport one size, h_stride = w_stride + // infact model stride is (stride, 1) + // but now setting stride all equal to 1 + // so h direction use indices select + if self.stride != 1 { + let h_dim = xs.dim(2)?; + let half_h = h_dim / 2; + let indices = Tensor::arange(0u32, half_h as u32, xs.device())?.affine(2.0, 0.0)?; + xs = xs.index_select(&indices, 2)?; + } + let xs = self.bn1.forward_t(&xs, false)?.relu()?; + let xs = self.conv2.forward(&xs)?; + let mut xs = self.bn2.forward_t(&xs, false)?; + if let Some(cut) = &self.shortcut { + let shortcut = cut.forward(&residual)?; + xs = xs.add(&shortcut)?; + } else { + xs = xs.add(&residual)?; + } + xs = xs.relu()?; + Ok(xs) + } +} + +pub struct FCM { + conv1: Conv2d, + bn1: BatchNorm, + layer1: Vec, + layer2: Vec, + conv2: Conv2d, + bn2: BatchNorm, + pub out_channels: usize, +} + +impl FCM { + pub fn new( + vb: VarBuilder, + num_blocks: &[usize], + m_channels: usize, + feat_dim: usize, + ) -> Result { + let conv1 = get_conv2d(vb.pp("conv1"), 1, m_channels, 3, 1, 1, 1, 1, false)?; + let bn1 = get_batch_norm(vb.pp("bn1"), 1e-5, m_channels, true)?; + let layer1_num_blocks = num_blocks[0] - 1; + let strides: Vec = [2usize] + .into_iter() + .chain([1usize].into_iter().cycle().take(layer1_num_blocks)) + .collect(); + let mut layer1 = vec![]; + let vb_layer1 = vb.pp("layer1"); + for (i, stride) in strides.iter().enumerate() { + let layer = BasicResBlock::new(vb_layer1.pp(i), m_channels, m_channels, *stride, 1)?; + layer1.push(layer); + } + let layer2_num_blocks = num_blocks[1] - 1; + let strides: Vec = [2usize] + .into_iter() + .chain([1usize].into_iter().cycle().take(layer2_num_blocks)) + .collect(); + let mut layer2 = vec![]; + let vb_layer2 = vb.pp("layer2"); + for (i, stride) in strides.iter().enumerate() { + let layer = BasicResBlock::new(vb_layer2.pp(i), m_channels, m_channels, *stride, 1)?; + layer2.push(layer); + } + let conv2 = get_conv2d(vb.pp("conv2"), m_channels, m_channels, 3, 1, 1, 1, 1, false)?; + let bn2 = get_batch_norm(vb.pp("bn2"), 1e-5, m_channels, true)?; + let out_channels = m_channels * (feat_dim / 8); + Ok(Self { + conv1, + bn1, + layer1, + layer2, + conv2, + bn2, + out_channels, + }) + } + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.unsqueeze(1)?; + let xs = self.conv1.forward(&xs)?; + let mut xs = self.bn1.forward_t(&xs, false)?.relu()?; + for layer in &self.layer1 { + xs = layer.forward(&xs)?; + } + for layer in &self.layer2 { + xs = layer.forward(&xs)?; + } + xs = self.conv2.forward(&xs)?; + let h_dim = xs.dim(2)?; + let half_h = h_dim / 2; + let indices = Tensor::arange(0u32, half_h as u32, xs.device())?.affine(2.0, 0.0)?; + xs = xs.index_select(&indices, 2)?; + xs = self.bn2.forward_t(&xs, false)?.relu()?; + let (bs, c, h, dim) = xs.dims4()?; + xs = xs.reshape((bs, c * h, dim))?; + Ok(xs) + } +} + +pub struct TDNNLayer { + linear: Conv1d, + nonlinear: BatchNorm, +} + +impl TDNNLayer { + pub fn new( + vb: VarBuilder, + in_c: usize, + out_c: usize, + ks: usize, + stride: usize, + dilation: usize, + bias: bool, + ) -> Result { + let padding = (ks - 1) / 2 * dilation; + let linear = get_conv1d( + vb.pp("linear"), + in_c, + out_c, + ks, + padding, + stride, + dilation, + 1, + bias, + )?; + let nonlinear = get_batch_norm(vb.pp("nonlinear.batchnorm"), 1e-5, out_c, true)?; + Ok(Self { linear, nonlinear }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = self.linear.forward(xs)?; + let xs = self.nonlinear.forward_t(&xs, false)?.relu()?; + Ok(xs) + } +} + +pub struct CAMLayer { + linear_local: Conv1d, + linear1: Conv1d, + linear2: Conv1d, +} + +impl CAMLayer { + pub fn new( + vb: VarBuilder, + bn_c: usize, + out_c: usize, + ks: usize, + stride: usize, + padding: usize, + dilation: usize, + bias: bool, + reduction: usize, + ) -> Result { + let linear_local = get_conv1d( + vb.pp("linear_local"), + bn_c, + out_c, + ks, + padding, + stride, + dilation, + 1, + bias, + )?; + let linear1 = get_conv1d( + vb.pp("linear1"), + bn_c, + bn_c / reduction, + 1, + 0, + 1, + 1, + 1, + true, + )?; + let linear2 = get_conv1d( + vb.pp("linear2"), + bn_c / reduction, + out_c, + 1, + 0, + 1, + 1, + 1, + true, + )?; + Ok(Self { + linear_local, + linear1, + linear2, + }) + } + + pub fn seg_pooling(&self, xs: &Tensor, seg_len: usize, stype: &str) -> Result { + let x_dim = xs.dim(2)?; + let seg = pool1d(xs, seg_len, true, stype)?; + let (bs, c, dim) = seg.dims3()?; + let seg = seg + .unsqueeze(D::Minus1)? + .expand((bs, c, dim, seg_len))? + .reshape((bs, c, ()))?; + let seg = seg.narrow(D::Minus1, 0, x_dim)?; + Ok(seg) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let y = self.linear_local.forward(xs)?; + let x_pool = self.seg_pooling(xs, 100, "avg")?; + let context = xs.mean_keepdim(D::Minus1)?.broadcast_add(&x_pool)?; + let context = self.linear1.forward(&context)?.relu()?; + let m = sigmoid(&self.linear2.forward(&context)?)?; + let res = y.mul(&m)?; + Ok(res) + } +} + +pub struct CAMDenseTDNNLayer { + nonlinear1: BatchNorm, + linear1: Conv1d, + nonlinear2: BatchNorm, + cam_layer: CAMLayer, +} + +impl CAMDenseTDNNLayer { + pub fn new( + vb: VarBuilder, + in_c: usize, + out_c: usize, + bn_c: usize, + ks: usize, + stride: usize, + dilation: usize, + bias: bool, + ) -> Result { + let padding = (ks - 1) / 2 * dilation; + let nonlinear1 = get_batch_norm(vb.pp("nonlinear1.batchnorm"), 1e-5, in_c, true)?; + let linear1 = get_conv1d(vb.pp("linear1"), in_c, bn_c, 1, 0, 1, 1, 1, false)?; + let nonlinear2 = get_batch_norm(vb.pp("nonlinear2.batchnorm"), 1e-5, bn_c, true)?; + let cam_layer = CAMLayer::new( + vb.pp("cam_layer"), + bn_c, + out_c, + ks, + stride, + padding, + dilation, + bias, + 2, + )?; + Ok(Self { + nonlinear1, + linear1, + nonlinear2, + cam_layer, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = self.nonlinear1.forward_t(xs, false)?.relu()?; + let xs = self.linear1.forward(&xs)?; + let xs = self.nonlinear2.forward_t(&xs, false)?.relu()?; + let xs = self.cam_layer.forward(&xs)?; + Ok(xs) + } +} + +pub struct CAMDenseTDNNBlock { + tdnns: Vec, +} + +impl CAMDenseTDNNBlock { + pub fn new( + vb: VarBuilder, + num_layers: usize, + in_c: usize, + out_c: usize, + bn_c: usize, + ks: usize, + stride: usize, + dilation: usize, + bias: bool, + ) -> Result { + let mut tdnns = vec![]; + for i in 0..num_layers { + let layer = CAMDenseTDNNLayer::new( + vb.pp(format!("tdnnd{}", i + 1)), + in_c + i * out_c, + out_c, + bn_c, + ks, + stride, + dilation, + bias, + )?; + tdnns.push(layer); + } + Ok(Self { tdnns }) + } + pub fn forward(&self, xs: &Tensor) -> Result { + let mut xs = xs.clone(); + for layer in &self.tdnns { + let layer_out = layer.forward(&xs)?; + xs = Tensor::cat(&[&xs, &layer_out], 1)?; + } + Ok(xs) + } +} + +pub struct TransitLayer { + nonlinear: BatchNorm, + linear: Conv1d, +} + +impl TransitLayer { + pub fn new(vb: VarBuilder, in_c: usize, out_c: usize, bias: bool) -> Result { + let nonlinear = get_batch_norm(vb.pp("nonlinear.batchnorm"), 1e-5, in_c, true)?; + let linear = get_conv1d(vb.pp("linear"), in_c, out_c, 1, 0, 1, 1, 1, bias)?; + Ok(Self { nonlinear, linear }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = self.nonlinear.forward_t(xs, false)?.relu()?; + let xs = self.linear.forward(&xs)?; + Ok(xs) + } +} + +pub struct DenseLayer { + linear: Conv1d, + nonlinear: BatchNorm, // only batch norm, no relu +} + +impl DenseLayer { + pub fn new(vb: VarBuilder, in_c: usize, out_c: usize, bias: bool) -> Result { + let linear = get_conv1d(vb.pp("linear"), in_c, out_c, 1, 0, 1, 1, 1, bias)?; + let nonlinear = get_batch_norm(vb.pp("nonlinear.batchnorm"), 1e-5, out_c, false)?; + Ok(Self { linear, nonlinear }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = if xs.rank() == 2 { + self.linear + .forward(&xs.unsqueeze(D::Minus1)?)? + .squeeze(D::Minus1)? + } else { + self.linear.forward(&xs)? + }; + let xs = self.nonlinear.forward_t(&xs, false)?; + Ok(xs) + } +} + +pub struct XVector { + tdnn: TDNNLayer, + blocks: Vec, + transits: Vec, + out_nonlinear: BatchNorm, + dense: DenseLayer, +} + +impl XVector { + pub fn new( + vb: VarBuilder, + channels: usize, + init_channels: usize, + growth_rate: usize, + bn_size: usize, + embedding_size: usize, + ) -> Result { + let tdnn = TDNNLayer::new(vb.pp("tdnn"), channels, init_channels, 5, 2, 1, false)?; + let mut channels = init_channels; + let mut blocks = vec![]; + let mut transits = vec![]; + let params = vec![(12, 3, 1), (24, 3, 2), (16, 3, 2)]; + for (i, (num_layers, ks, dilation)) in params.iter().enumerate() { + let block = CAMDenseTDNNBlock::new( + vb.pp(format!("block{}", i + 1)), + *num_layers, + channels, + growth_rate, + bn_size * growth_rate, + *ks, + 1, + *dilation, + false, + )?; + blocks.push(block); + channels = channels + num_layers * growth_rate; + let transit = TransitLayer::new( + vb.pp(format!("transit{}", i + 1)), + channels, + channels / 2, + false, + )?; + transits.push(transit); + channels /= 2; + } + let out_nonlinear = get_batch_norm(vb.pp("out_nonlinear.batchnorm"), 1e-5, channels, true)?; + let dense = DenseLayer::new(vb.pp("dense"), channels * 2, embedding_size, false)?; + Ok(Self { + tdnn, + blocks, + transits, + out_nonlinear, + dense, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let mut xs = self.tdnn.forward(xs)?; + for i in 0..3 { + let block = &self.blocks[i]; + xs = block.forward(&xs)?; + let transit = &self.transits[i]; + xs = transit.forward(&xs)?; + } + xs = self.out_nonlinear.forward_t(&xs, false)?.relu()?; + xs = statistics_pooling(&xs, D::Minus1, false)?; + xs = self.dense.forward(&xs)?; + Ok(xs) + } +} + +pub struct CAMPPlus { + head: FCM, + xvector: XVector, +} + +impl CAMPPlus { + pub fn new( + vb: VarBuilder, + feat_dim: usize, + embedding_size: usize, + growth_rate: usize, + bn_size: usize, + init_channels: usize, + ) -> Result { + let head = FCM::new(vb.pp("head"), &[2, 2], 32, feat_dim)?; + let channels = head.out_channels; + let xvector = XVector::new( + vb.pp("xvector"), + channels, + init_channels, + growth_rate, + bn_size, + embedding_size, + )?; + Ok(Self { head, xvector }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = xs.permute((0, 2, 1))?; + let xs = self.head.forward(&xs)?; + let xs = self.xvector.forward(&xs)?; + Ok(xs) + } +} diff --git a/src/models/common/mod.rs b/src/models/common/mod.rs index 9670ac1..2b4f46f 100644 --- a/src/models/common/mod.rs +++ b/src/models/common/mod.rs @@ -1,15 +1,15 @@ use anyhow::{Result, anyhow}; -use candle_core::{D, Tensor}; +use candle_core::{D, IndexOp, Tensor}; use candle_nn::{ - Activation, BatchNorm, BatchNormConfig, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig, Embedding, - LayerNorm, LayerNormConfig, Linear, Module, RmsNorm, VarBuilder, batch_norm, conv1d, - conv1d_no_bias, conv2d, conv2d_no_bias, embedding, layer_norm, linear_b, linear_no_bias, - rms_norm, + Activation, BatchNorm, BatchNormConfig, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig, + ConvTranspose1d, ConvTranspose1dConfig, Embedding, LayerNorm, LayerNormConfig, Linear, Module, + ModuleT, RmsNorm, VarBuilder, batch_norm, conv1d, conv1d_no_bias, conv2d, conv2d_no_bias, + embedding, layer_norm, linear_b, linear_no_bias, ops::sigmoid, rms_norm, }; use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator}; use crate::{ - position_embed::rope::{RoPE, apply_rotary_pos_emb}, + position_embed::rope::{RoPE, apply_rotary_pos_emb, apply_rotary_pos_emb_roformer}, utils::tensor_utils::{prepare_causal_attention_mask, repeat_kv}, }; @@ -28,10 +28,16 @@ impl GateUpDownMLP { intermediate_size: usize, act_fn: Activation, bias: bool, + gate_pp_name: Option<&str>, + up_pp_name: Option<&str>, + down_pp_name: Option<&str>, ) -> Result { - let gate_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp("gate_proj"))?; - let up_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp("up_proj"))?; - let down_proj = linear_b(intermediate_size, hidden_size, bias, vb.pp("down_proj"))?; + let gate_pp_name = gate_pp_name.unwrap_or("gate_proj"); + let up_pp_name = up_pp_name.unwrap_or("up_proj"); + let down_pp_name = down_pp_name.unwrap_or("down_proj"); + let gate_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp(gate_pp_name))?; + let up_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp(up_pp_name))?; + let down_proj = linear_b(intermediate_size, hidden_size, bias, vb.pp(down_pp_name))?; Ok(Self { gate_proj, up_proj, @@ -87,7 +93,6 @@ impl TwoLinearMLP { } #[derive(Debug, Clone)] -// pub struct AttentionNobias { pub struct NaiveAttention { q_proj: Linear, k_proj: Linear, @@ -257,6 +262,154 @@ impl NaiveAttention { } } +#[derive(Debug, Clone)] +pub struct QKVCatAttention { + qkv_proj: Linear, + o_proj: Linear, + num_heads: usize, + head_dim: usize, + middle_size: usize, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl QKVCatAttention { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + num_attention_heads: usize, + head_dim: Option, + bias: bool, + qkv_proj_pp_name: Option<&str>, + o_proj_pp_name: Option<&str>, + ) -> Result { + let head_dim = match head_dim { + None => hidden_size / num_attention_heads, + Some(dim) => dim, + }; + let qkv_proj_pp_name = qkv_proj_pp_name.unwrap_or("wqkv"); + let o_proj_pp_name = o_proj_pp_name.unwrap_or("o_proj"); + let qkv_proj = linear_b( + hidden_size, + 3 * num_attention_heads * head_dim, + bias, + vb.pp(qkv_proj_pp_name), + )?; + let o_proj = linear_b( + num_attention_heads * head_dim, + hidden_size, + bias, + vb.pp(o_proj_pp_name), + )?; + + Ok(Self { + qkv_proj, + o_proj, + num_heads: num_attention_heads, + head_dim, + middle_size: num_attention_heads * head_dim, + kv_cache: None, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + cos: Option<&Tensor>, + sin: Option<&Tensor>, + attention_mask: Option<&Tensor>, + tof32: bool, + use_roformer: bool, + ) -> Result { + let (b, q_len, _) = xs.dims3()?; + // (3, B, n_head, seq_len, head_dim) + let qkv = self + .qkv_proj + .forward(xs)? + .reshape((b, q_len, 3, self.num_heads, ()))? + .permute((2, 0, 3, 1, 4))? + .contiguous()?; + let query_states = qkv.i(0)?.contiguous()?; + let key_states = qkv.i(1)?.contiguous()?; + let value_states = qkv.i(2)?.contiguous()?; + let (query_states, key_states) = if let Some(cos) = cos + && let Some(sin) = sin + { + if use_roformer { + apply_rotary_pos_emb_roformer(&query_states, &key_states, cos, sin, tof32)? + } else { + apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)? + } + } else { + (query_states, key_states) + }; + + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_output = eager_attention_forward( + &query_states, + &key_states, + &value_states, + None, + attention_mask, + scale, + )?; + let attn_output = attn_output.reshape((b, q_len, self.middle_size))?; + let attn_output = attn_output.apply(&self.o_proj)?; + Ok(attn_output) + } + + pub fn forward_with_cache( + &mut self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + tof32: bool, + use_roformer: bool, + ) -> Result { + let (b, q_len, _) = xs.dims3()?; + let qkv = self + .qkv_proj + .forward(xs)? + .reshape((b, q_len, 3, self.num_heads, ()))? + .permute((2, 0, 3, 1, 4))? + .contiguous()?; + let query_states = qkv.i(0)?.contiguous()?; + let key_states = qkv.i(1)?.contiguous()?; + let value_states = qkv.i(2)?.contiguous()?; + let (query_states, key_states) = if use_roformer { + apply_rotary_pos_emb_roformer(&query_states, &key_states, cos, sin, tof32)? + } else { + apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)? + }; + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + + self.kv_cache = Some((key_states.clone(), value_states.clone())); + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attn_output = eager_attention_forward( + &query_states, + &key_states, + &value_states, + None, + attention_mask, + scale, + )?; + let attn_output = attn_output.reshape((b, q_len, self.middle_size))?; + let attn_output = attn_output.apply(&self.o_proj)?; + Ok(attn_output) + } + + pub fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + pub struct NaiveAttnTwoLinearMLPBlock { self_attn: NaiveAttention, mlp: TwoLinearMLP, @@ -390,6 +543,9 @@ impl NaiveAttnGateUpDownMLPBlock { intermediate_size, hidden_act, mlp_bias, + None, + None, + None, )?; let input_layernorm = rms_norm(hidden_size, norm_eps, vb.pp(input_norm_pp_name))?; let post_attention_layernorm = rms_norm(hidden_size, norm_eps, vb.pp(post_norm_pp_name))?; @@ -543,11 +699,11 @@ pub fn get_layer_norm(vb: VarBuilder, eps: f64, dim: usize) -> Result Ok(norm) } -pub fn get_batch_norm(vb: VarBuilder, eps: f64, dim: usize) -> Result { +pub fn get_batch_norm(vb: VarBuilder, eps: f64, dim: usize, affine: bool) -> Result { let bn_config = BatchNormConfig { eps, remove_mean: true, - affine: true, + affine, momentum: 0.1, }; let norm = batch_norm(dim, bn_config, vb)?; @@ -850,3 +1006,164 @@ pub fn conv1d_group_parallel(xs: &Tensor, conv1d: &Conv1d) -> Result { } } } + +pub struct GLU { + dim: usize, +} + +impl GLU { + pub fn new(dim: usize) -> Result { + Ok(Self { dim }) + } + pub fn forward(&self, xs: &Tensor) -> Result { + let half_dim = xs.dim(self.dim)? / 2; + let a = xs.narrow(self.dim, 0, half_dim)?; + let b = xs.narrow(self.dim, half_dim, half_dim)?; + let b = sigmoid(&b)?; + let xs = a.mul(&b)?; + Ok(xs) + } +} + +pub struct WNConv1d { + conv: Conv1d, +} +impl WNConv1d { + pub fn new( + vb: VarBuilder, + in_c: usize, + out_c: usize, + kernel_size: usize, + dilation: usize, + padding: usize, + groups: usize, + stride: usize, + bias: bool, + ) -> Result { + let in_c = in_c / groups; + let weight_g = vb.get((out_c, 1, 1), "weight_g")?; + let weight_v = vb.get((out_c, in_c, kernel_size), "weight_v")?; + // let bias = vb.get(out_c, "bias").ok(); + let bias = if bias { + vb.get(out_c, "bias").ok() + } else { + None + }; + let weight_norm = weight_v.sqr()?.sum_keepdim(1)?.sum_keepdim(2)?.sqrt()?; + let normalized_weight = weight_v.broadcast_div(&weight_norm)?; + let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?; + let cfg = Conv1dConfig { + padding, + stride, + dilation, + groups, + cudnn_fwd_algo: None, + }; + let conv = Conv1d::new(scaled_weight, bias, cfg); + Ok(Self { conv }) + } + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.conv.forward(x)?; + Ok(x) + } +} + +pub struct WNConvTranspose1d { + conv_transpose: ConvTranspose1d, +} + +impl WNConvTranspose1d { + pub fn new( + vb: VarBuilder, + in_c: usize, + out_c: usize, + dilation: usize, + kernel_size: usize, + padding: usize, + output_padding: usize, + groups: usize, + stride: usize, + ) -> Result { + let in_c = in_c / groups; + let weight_g = vb.get((in_c, 1, 1), "weight_g")?; + let weight_v = vb.get((in_c, out_c, kernel_size), "weight_v")?; + let bias = vb.get(out_c, "bias").ok(); + let weight_norm = weight_v.sqr()?.sum_keepdim(1)?.sum_keepdim(2)?.sqrt()?; + let normalized_weight = weight_v.broadcast_div(&weight_norm)?; + let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?; + let config = ConvTranspose1dConfig { + padding: padding, + output_padding: output_padding, + stride, + dilation, + groups, + }; + let conv_transpose = ConvTranspose1d::new(scaled_weight, bias, config); + Ok(Self { conv_transpose }) + } + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.conv_transpose.forward(x)?; + Ok(x) + } +} + +pub struct Conv2dWithBN { + conv_0: Conv2d, + bn_1: BatchNorm, + bn_with_relu: bool, +} + +impl Conv2dWithBN { + pub fn new( + vb: VarBuilder, + in_c: usize, + out_c: usize, + ks: usize, + padding: usize, + stride: usize, + bias: bool, + bn_with_relu: bool, + ) -> Result { + let conv_0 = get_conv2d(vb.pp("0"), in_c, out_c, ks, padding, stride, 1, 1, bias)?; + let bn_1 = get_batch_norm(vb.pp("1"), 1e-5, out_c, true)?; + Ok(Self { + conv_0, + bn_1, + bn_with_relu, + }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.conv_0.forward(x)?; + let mut x = self.bn_1.forward_t(&x, false)?; + if self.bn_with_relu { + x = x.relu()?; + } + Ok(x) + } +} + +pub struct WNLinear { + linear: Linear, +} +impl WNLinear { + pub fn new(vb: VarBuilder, in_dim: usize, out_dim: usize, bias: bool) -> Result { + let weight_g = vb.get((out_dim, 1), "weight_g")?; + let weight_v = vb.get((out_dim, in_dim), "weight_v")?; + + let bias = if bias { + vb.get(out_dim, "bias").ok() + } else { + None + }; + let weight_norm = weight_v.sqr()?.sum_keepdim(0)?.sqrt()?.affine(1.0, 1e-8)?; + let normalized_weight = weight_v.broadcast_div(&weight_norm)?; + let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?; + let linear = Linear::new(scaled_weight, bias); + Ok(Self { linear }) + } + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.linear.forward(x)?; + Ok(x) + } +} diff --git a/src/models/deepseek_ocr/model.rs b/src/models/deepseek_ocr/model.rs index 02c6560..493d214 100644 --- a/src/models/deepseek_ocr/model.rs +++ b/src/models/deepseek_ocr/model.rs @@ -875,6 +875,9 @@ impl DeepseekV2MoE { config.moe_intermediate_size, Activation::Silu, false, + None, + None, + None, )?; experts.push(mlp); } @@ -885,6 +888,9 @@ impl DeepseekV2MoE { config.moe_intermediate_size * config.n_shared_experts, Activation::Silu, false, + None, + None, + None, )?; Ok(Self { // num_experts_per_tok: config.num_experts_per_tok, @@ -998,6 +1004,9 @@ impl DeepseekV2DecoderLayer { config.intermediate_size, Activation::Silu, false, + None, + None, + None, )?) }; let input_layernorm = rms_norm( diff --git a/src/models/feature_extractor/mod.rs b/src/models/feature_extractor/mod.rs new file mode 100644 index 0000000..5fb1e5a --- /dev/null +++ b/src/models/feature_extractor/mod.rs @@ -0,0 +1 @@ +pub mod seamless_m4t_feature_extractor; \ No newline at end of file diff --git a/src/models/feature_extractor/seamless_m4t_feature_extractor.rs b/src/models/feature_extractor/seamless_m4t_feature_extractor.rs new file mode 100644 index 0000000..aa371cb --- /dev/null +++ b/src/models/feature_extractor/seamless_m4t_feature_extractor.rs @@ -0,0 +1,129 @@ +use anyhow::Result; +use candle_core::{D, Device, Tensor}; + +use crate::utils::{ + audio_utils::{create_povey_window, mel_filter_bank, spectrogram}, + tensor_utils::{PaddingSide, z_score_normalize}, +}; + +pub struct SeamlessM4TFeatureExtractor { + feature_size: usize, + num_mel_bins: usize, + padding_side: PaddingSide, + padding_value: f32, + sampling_rate: usize, + stride: usize, + mel_filters: Tensor, + window: Tensor, +} + +impl SeamlessM4TFeatureExtractor { + pub fn new( + feature_size: usize, + num_mel_bins: usize, + padding_side: PaddingSide, + padding_value: f32, + sampling_rate: usize, + stride: usize, + device: &Device, + ) -> Result { + let mel_filters = mel_filter_bank( + 257, + num_mel_bins, + 20.0, + (sampling_rate / 2) as f32, + sampling_rate as f32, + None, + crate::utils::audio_utils::MelScale::Kaldi, + true, + device, + )?; + let window = create_povey_window(400, candle_core::DType::F32, device)?; + Ok(Self { + feature_size, + num_mel_bins, + padding_side, + padding_value, + sampling_rate, + stride, + mel_filters, + window, + }) + } + + pub fn call( + &self, + raw_speech: &Tensor, + sampling_rate: usize, + do_normalize_per_mel_bins: bool, + return_attention_mask: bool, + ) -> Result<(Tensor, Option)> { + // raw_speech: 重采样后的音频，shape: (bs, raw_len) + // sampling_rate: 音频采样率，验证是否与模型的预处理采样率一致 + if sampling_rate != self.sampling_rate { + return Err(anyhow::anyhow!( + "The model feature extractor was trained sampling rate {} not equal to audio sample rate {}", + self.sampling_rate, + sampling_rate + )); + } + let waveform = raw_speech.affine(32768.0, 0.0)?; + // println!("waveform: {}", waveform); + // println!("self.mel_filters: {}", self.mel_filters); + // println!("self.window: {}", self.window); + let mut features = spectrogram( + &waveform, + &self.window, + 400, + 160, + 512, + Some(2.0), + false, + 0.97, + Some(&self.mel_filters), + Some("log"), + 1.192092955078125e-07, + true, + )? + .transpose(D::Minus1, D::Minus2)?; + if do_normalize_per_mel_bins { + features = z_score_normalize(&features, 1)?; + } + let n_frame = features.dim(1)?; + let mask_1 = n_frame / self.stride; + let pad_len = n_frame % self.stride; + if pad_len > 0 { + let pad = Tensor::new(self.padding_value, features.device())?.broadcast_as(( + 1, + pad_len, + self.num_mel_bins, + ))?; + match self.padding_side { + PaddingSide::Left => features = Tensor::cat(&[pad, features], 1)?, + PaddingSide::Right => features = Tensor::cat(&[features, pad], 1)?, + } + } + let (bs, num_frames, dim) = features.dims3()?; + let n_frames_stride = num_frames / self.stride; + let features = features.reshape((bs, n_frames_stride, dim * self.stride))?; + let mask_0 = n_frames_stride - mask_1; + let mask = if return_attention_mask { + let mut mask = Tensor::new(1u32, features.device())?.broadcast_as((1, mask_1))?; + if mask_0 > 0 { + let mask_pad = Tensor::new(0u32, features.device())?.broadcast_as((1, mask_0))?; + match self.padding_side { + PaddingSide::Left => { + mask = Tensor::cat(&[mask_pad, mask], D::Minus1)?; + } + PaddingSide::Right => { + mask = Tensor::cat(&[mask, mask_pad], D::Minus1)?; + } + } + } + Some(mask) + } else { + None + }; + Ok((features, mask)) + } +} diff --git a/src/models/fun_asr_nano/model.rs b/src/models/fun_asr_nano/model.rs index d3f971f..755ecff 100644 --- a/src/models/fun_asr_nano/model.rs +++ b/src/models/fun_asr_nano/model.rs @@ -11,7 +11,7 @@ use crate::{ qwen3::{config::Qwen3Config, model::Qwen3Model}, }, position_embed::sinusoidal_pe::SinusoidalPositionEncoderCat, - utils::tensor_utils::{get_equal_mask, mask_filled, masked_scatter_dim0}, + utils::tensor_utils::{attn_masked_fill, get_equal_mask, masked_scatter_dim0}, }; pub struct MultiHeadedAttentionSANM { @@ -124,9 +124,9 @@ impl MultiHeadedAttentionSANM { }; // mask: rank = 2 let mask = get_equal_mask(&mask, 0)?; - let scores = mask_filled(scores, &mask, f32::NEG_INFINITY)?; + let scores = attn_masked_fill(scores, &mask, f32::NEG_INFINITY)?; let attn = softmax_last_dim(&scores)?; - mask_filled(&attn, &mask, 0.0)? + attn_masked_fill(&attn, &mask, 0.0)? } else { softmax_last_dim(scores)? }; diff --git a/src/models/glm_asr_nano/processor.rs b/src/models/glm_asr_nano/processor.rs index 5ba02da..52783fa 100644 --- a/src/models/glm_asr_nano/processor.rs +++ b/src/models/glm_asr_nano/processor.rs @@ -9,8 +9,9 @@ use crate::{ utils::{ audio_utils::{ apply_stft, create_hann_window, extract_audios, extract_frames, mel_filter_bank, + torch_stft, }, - tensor_utils::{pad_reflect_last_dim, split_tensor}, + tensor_utils::{log10, pad_reflect_last_dim, split_tensor}, }, }; @@ -87,19 +88,21 @@ impl GlmAsrNanoProcessor { let waveform = pad_reflect_last_dim(waveform, (pad, pad))?; let (_, samples) = waveform.dims2()?; - // 计算输出维度 + // // (bs, n_frames, n_fft) + // let frames = extract_frames(&waveform, self.n_fft, self.hop_length)?; + // // 应用汉明窗口 + // let result = frames.broadcast_mul(&self.window)?; + // // 傅立叶变换 + // let magnitudes = apply_stft(&result)?.transpose(D::Minus1, D::Minus2)?; + let magnitudes = torch_stft(&waveform, self.n_fft, self.hop_length, &self.window)? + .transpose(D::Minus1, D::Minus2)?; let n_frames = (samples - self.n_fft) / self.hop_length + 1; - // (bs, n_frames, n_fft) - let frames = extract_frames(&waveform, self.n_fft, self.hop_length)?; - // 应用汉明窗口 - let result = frames.broadcast_mul(&self.window)?; - // 傅立叶变换 - let magnitudes = apply_stft(&result)?.transpose(D::Minus1, D::Minus2)?; let magnitudes = magnitudes.narrow(D::Minus1, 0, n_frames - 1)?; let mel_spec = self.mel_filters.broadcast_matmul(&magnitudes)?; let mel_spec = mel_spec.clamp(1e-10f32, f32::INFINITY)?; - let ln_spec = mel_spec.log()?; - let log10_spec = ln_spec.broadcast_div(&Tensor::new(f32::ln(10.0), mel_spec.device())?)?; + // let ln_spec = mel_spec.log()?; + // let log10_spec = ln_spec.broadcast_div(&Tensor::new(f32::ln(10.0), mel_spec.device())?)?; + let log10_spec = log10(&mel_spec)?; let max_val = log10_spec.max_all()?.affine(1.0, -8.0)?; let log10_spec = log10_spec.broadcast_maximum(&max_val)?; let log_spec = log10_spec.affine(1.0, 4.0)?.affine(1.0 / 4.0, 0.0)?; diff --git a/src/models/hunyuan_ocr/model.rs b/src/models/hunyuan_ocr/model.rs index 4df4ffb..c9d4b60 100644 --- a/src/models/hunyuan_ocr/model.rs +++ b/src/models/hunyuan_ocr/model.rs @@ -406,6 +406,9 @@ impl HunYuanVLDecoderLayer { config.intermediate_size, config.hidden_act, false, + None, + None, + None, )?; let input_layernorm = rms_norm( config.hidden_size, diff --git a/src/models/index_tts2/config.rs b/src/models/index_tts2/config.rs new file mode 100644 index 0000000..e6affb6 --- /dev/null +++ b/src/models/index_tts2/config.rs @@ -0,0 +1,228 @@ +use serde::{Deserialize, Deserializer}; + +use crate::models::mask_gct::config::SemanticCodec; + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct IndexTTS2Config { + pub dataset: Dataset, + pub gpt: Gpt, + pub semantic_codec: SemanticCodec, + pub s2mel: S2MelConfig, + pub gpt_checkpoint: String, + pub w2v_stat: String, + pub s2mel_checkpoint: String, + pub emo_matrix: String, + pub spk_matrix: String, + pub emo_num: Vec, + pub qwen_emo_path: String, + pub vocoder: Vocoder, + pub version: f32, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Dataset { + pub bpe_model: String, + pub sample_rate: usize, + pub squeeze: bool, + pub mel: Mel, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Mel { + pub sample_rate: usize, + pub n_fft: usize, + pub hop_length: usize, + pub win_length: usize, + pub n_mels: usize, + pub mel_fmin: usize, + pub normalize: bool, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Gpt { + pub model_dim: usize, + pub max_mel_tokens: usize, + pub max_text_tokens: usize, + pub heads: usize, + pub use_mel_codes_as_input: bool, + pub mel_length_compression: usize, + pub layers: usize, + pub number_text_tokens: usize, + pub number_mel_codes: usize, + pub start_mel_token: usize, + pub stop_mel_token: usize, + pub start_text_token: usize, + pub stop_text_token: usize, + pub train_solo_embeddings: bool, + pub condition_type: String, + pub condition_module: ConditionModule, + pub emo_condition_module: EmoConditionModule, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct ConditionModule { + pub output_size: usize, + pub linear_units: usize, + pub attention_heads: usize, + pub num_blocks: usize, + pub input_layer: String, + pub perceiver_mult: usize, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct EmoConditionModule { + pub output_size: usize, + pub linear_units: usize, + pub attention_heads: usize, + pub num_blocks: usize, + pub input_layer: String, + pub perceiver_mult: usize, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct S2MelConfig { + pub preprocess_params: PreprocessParams, + pub dit_type: String, + pub reg_loss_type: String, + pub style_encoder: StyleEncoder, + pub length_regulator: LengthRegulator, + #[serde(rename = "DiT")] + pub di_t: DiTConfig, + pub wavenet: WavenetConfig, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct PreprocessParams { + pub sr: usize, + pub spect_params: SpectParams, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct SpectParams { + pub n_fft: usize, + pub win_length: usize, + pub hop_length: usize, + pub n_mels: usize, + pub fmin: usize, + #[serde(deserialize_with = "deserialize_optional_fmax")] + pub fmax: Option, +} + +fn deserialize_optional_fmax<'de, D>(deserializer: D) -> Result, D::Error> +where D: Deserializer<'de> { + let opt: Option = Option::deserialize(deserializer)?; + match opt { + Some(serde_json::Value::String(s)) if s == "None" || s == "null" => Ok(None), + Some(serde_json::Value::Number(n)) => { + if let Some(n) = n.as_u64() { + Ok(Some(n as usize)) + } else { + Err(serde::de::Error::custom("Expected positive integer")) + } + } + Some(_) => Err(serde::de::Error::custom("Expected number or 'None' string")), + None => Ok(None), + } +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct StyleEncoder { + pub dim: usize, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct LengthRegulator { + pub channels: usize, + pub is_discrete: bool, + pub in_channels: usize, + pub content_codebook_size: usize, + pub sampling_ratios: Vec, + pub vector_quantize: bool, + pub n_codebooks: usize, + pub quantizer_dropout: f32, + pub f0_condition: bool, + pub n_f0_bins: usize, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct DiTConfig { + pub hidden_dim: usize, + pub num_heads: usize, + pub depth: usize, + pub class_dropout_prob: f32, + pub block_size: usize, + pub in_channels: usize, + pub style_condition: bool, + pub final_layer_type: String, + pub target: String, + pub content_dim: usize, + pub content_codebook_size: usize, + pub content_type: String, + pub f0_condition: bool, + pub n_f0_bins: usize, + pub content_codebooks: usize, + pub is_causal: bool, + pub long_skip_connection: bool, + pub zero_prompt_speech_token: bool, + pub time_as_token: bool, + pub style_as_token: bool, + pub uvit_skip_connection: bool, + pub add_resblock_in_transformer: bool, +} + +pub struct DiTModelArgs { + pub block_size: usize, + pub vocab_size: usize, + pub n_layer: usize, + pub n_head: usize, + pub dim: usize, + pub intermediate_size: usize, + pub n_local_heads: usize, + pub head_dim: usize, + pub rope_base: f32, + pub norm_eps: f64, + pub has_cross_attention: bool, + pub context_dim: usize, + pub uvit_skip_connection: bool, + pub time_as_token: bool, +} + +impl DiTModelArgs { + pub fn new_from_dit_config(config: &DiTConfig) -> Self { + let hidden_dim = 4 * config.hidden_dim; + let n_hidden = 2 * hidden_dim / 3; + let intermediate_size = n_hidden + (256 - n_hidden % 256) % 256; + Self { + block_size: config.block_size, + vocab_size: 1024, + n_layer: config.depth, + n_head: config.num_heads, + dim: config.hidden_dim, + intermediate_size, + n_local_heads: config.num_heads, + head_dim: config.hidden_dim / config.num_heads, + rope_base: 10000.0, + norm_eps: 1e-5, + has_cross_attention: false, + context_dim: 0, + uvit_skip_connection: config.uvit_skip_connection, + time_as_token: config.uvit_skip_connection, + } + } +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct WavenetConfig { + pub hidden_dim: usize, + pub num_layers: usize, + pub kernel_size: usize, + pub dilation_rate: usize, + pub p_dropout: f32, + pub style_condition: bool, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct Vocoder { + pub r#type: String, + pub name: String, +} diff --git a/src/models/index_tts2/generate.rs b/src/models/index_tts2/generate.rs new file mode 100644 index 0000000..923e7a3 --- /dev/null +++ b/src/models/index_tts2/generate.rs @@ -0,0 +1,29 @@ +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; +use anyhow::Result; +use candle_core::{DType, Device}; + +use crate::{ + models::index_tts2::{config::IndexTTS2Config, processor::IndexTTS2Processor}, + utils::{get_default_save_dir, get_device, get_dtype}, +}; +pub struct IndexTTS2Generate { + processor: IndexTTS2Processor, + config: IndexTTS2Config, +} + +impl IndexTTS2Generate { + pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { + let config_path = path.to_string() + "/config.yaml"; + let save_dir = get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let config: IndexTTS2Config = serde_yaml::from_slice(&std::fs::read(config_path)?)?; + let device = get_device(device); + let dtype = get_dtype(dtype, "bf16"); + let processor = IndexTTS2Processor::new(path, &save_dir, &config, &device, dtype)?; + + Ok(Self { config, processor }) + } + pub fn generate(&mut self, mes: ChatCompletionParameters) -> Result<()> { + let _ = self.processor.process_info(&mes)?; + Ok(()) + } +} diff --git a/src/models/index_tts2/mod.rs b/src/models/index_tts2/mod.rs new file mode 100644 index 0000000..6634bcd --- /dev/null +++ b/src/models/index_tts2/mod.rs @@ -0,0 +1,5 @@ +pub mod config; +pub mod generate; +pub mod model; +pub mod processor; +pub mod utils; \ No newline at end of file diff --git a/src/models/index_tts2/model.rs b/src/models/index_tts2/model.rs new file mode 100644 index 0000000..3f9b692 --- /dev/null +++ b/src/models/index_tts2/model.rs @@ -0,0 +1,617 @@ +use anyhow::Result; +use candle_core::{D, Tensor}; +use candle_nn::{ + Conv1d, Embedding, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, linear, linear_b, + ops::sigmoid, rms_norm, +}; + +use crate::{ + models::{ + common::{ + GateUpDownMLP, QKVCatAttention, TwoLinearMLP, WNConv1d, WNLinear, get_conv1d, + get_layer_norm, + }, + index_tts2::config::{DiTModelArgs, S2MelConfig}, + }, + position_embed::rope::RoPE, + utils::tensor_utils::{pad_reflect_last_dim, split_tensor_with_size}, +}; +pub struct AdaptiveLayerNorm { + project_layer: Linear, + norm: RmsNorm, + d_model: usize, +} + +impl AdaptiveLayerNorm { + pub fn new(vb: VarBuilder, d_model: usize, eps: f64) -> Result { + let project_layer = linear(d_model, d_model * 2, vb.pp("project_layer"))?; + let norm = rms_norm(d_model, eps, vb.pp("norm"))?; + Ok(Self { + project_layer, + norm, + d_model, + }) + } + + pub fn forward(&self, xs: &Tensor, embedding: Option<&Tensor>) -> Result { + if let Some(embedding) = embedding { + let emb = self.project_layer.forward(embedding)?; + let emb_split = split_tensor_with_size(&emb, 2, D::Minus1)?; + let weight = &emb_split[0]; + let bias = &emb_split[1]; + Ok(self + .norm + .forward(xs)? + .broadcast_mul(weight)? + .broadcast_add(bias)?) + } else { + Ok(self.norm.forward(xs)?) + } + } +} + +pub struct DiTTransformerBlock { + attention: QKVCatAttention, + feed_forward: GateUpDownMLP, + ffn_norm: AdaptiveLayerNorm, + attention_norm: AdaptiveLayerNorm, + skip_in_linear: Option, + uvit_skip_connection: bool, + time_as_token: bool, +} + +impl DiTTransformerBlock { + pub fn new(vb: VarBuilder, config: &DiTModelArgs) -> Result { + let attention = QKVCatAttention::new( + vb.pp("attention"), + config.dim, + config.n_head, + Some(config.head_dim), + false, + Some("wqkv"), + Some("wo"), + )?; + + let feed_forward = GateUpDownMLP::new( + vb.pp("feed_forward"), + config.dim, + config.intermediate_size, + candle_nn::Activation::Silu, + false, + Some("w1"), + Some("w3"), + Some("w2"), + )?; + let ffn_norm = AdaptiveLayerNorm::new(vb.pp("ffn_norm"), config.dim, config.norm_eps)?; + let attention_norm = + AdaptiveLayerNorm::new(vb.pp("attention_norm"), config.dim, config.norm_eps)?; + let (skip_in_linear, uvit_skip_connection) = if config.uvit_skip_connection { + let skip_in_linear = linear(config.dim * 2, config.dim, vb.pp("skip_in_linear"))?; + (Some(skip_in_linear), config.uvit_skip_connection) + } else { + (None, false) + }; + Ok(Self { + attention, + feed_forward, + ffn_norm, + attention_norm, + skip_in_linear, + uvit_skip_connection, + time_as_token: config.time_as_token, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + c: &Tensor, + cos: Option<&Tensor>, + sin: Option<&Tensor>, + mask: Option<&Tensor>, + skip_in_x: Option<&Tensor>, + ) -> Result { + let c = if self.time_as_token { None } else { Some(c) }; + let mut xs = xs.clone(); + if self.uvit_skip_connection + && let Some(skip_in_x) = skip_in_x + && let Some(skip_in_linear) = &self.skip_in_linear + { + let cat = Tensor::cat(&[&xs, skip_in_x], D::Minus1)?; + xs = skip_in_linear.forward(&cat)?; + } + let xs = self + .attention + .forward( + &self.attention_norm.forward(&xs, c)?, + cos, + sin, + mask, + false, + true, + )? + .add(&xs)?; + let out = self + .feed_forward + .forward(&self.ffn_norm.forward(&xs, c)?)? + .add(&xs)?; + Ok(out) + } +} + +pub struct DiTTransformer { + layers: Vec, + norm: AdaptiveLayerNorm, + rope: RoPE, + uvit_skip_connection: bool, + layers_emit_skip: Vec, + layers_receive_skip: Vec, +} + +impl DiTTransformer { + pub fn new(vb: VarBuilder, config: &DiTModelArgs) -> Result { + let vb_layers = vb.pp("layers"); + let mut layers = vec![]; + for i in 0..config.n_layer { + let layer = DiTTransformerBlock::new(vb_layers.pp(i), config)?; + layers.push(layer); + } + let norm = AdaptiveLayerNorm::new(vb.pp("norm"), config.dim, config.norm_eps)?; + let rope = RoPE::new(config.dim, 10000.0, vb.device())?; + let mut layers_emit_skip: Vec = vec![]; + let mut layers_receive_skip: Vec = vec![]; + if config.uvit_skip_connection { + layers_emit_skip = (0..config.n_layer) + .filter(|&x| x < config.n_layer / 2) + .collect(); + layers_receive_skip = (0..config.n_layer) + .filter(|&x| x > config.n_layer / 2) + .collect(); + } + Ok(Self { + layers, + norm, + rope, + uvit_skip_connection: config.uvit_skip_connection, + layers_emit_skip, + layers_receive_skip, + }) + } + pub fn forward(&self, xs: &Tensor, c: &Tensor, mask: Option<&Tensor>) -> Result { + let (_, seq_len, _) = xs.dims3()?; + let (cos, sin) = self.rope.forward(0, seq_len, xs.device())?; + let mut skip_in_x_list = vec![]; + let mut xs = xs.clone(); + for (i, layer) in (&self.layers).iter().enumerate() { + let skip_in_x = if self.uvit_skip_connection && self.layers_receive_skip.contains(&i) { + skip_in_x_list.pop() + } else { + None + }; + xs = layer.forward(&xs, c, Some(&cos), Some(&sin), mask, skip_in_x.as_ref())?; + if self.uvit_skip_connection && self.layers_emit_skip.contains(&i) { + skip_in_x_list.push(xs.clone()); + } + } + xs = self.norm.forward(&xs, Some(c))?; + Ok(xs) + } +} + +pub struct TimestepEmbedder { + mlp: TwoLinearMLP, + freqs: Tensor, + scale: f64, + frequency_embedding_size: usize, +} + +impl TimestepEmbedder { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + frequency_embedding_size: usize, + ) -> Result { + let mlp = TwoLinearMLP::new( + vb.pp("mlp"), + frequency_embedding_size, + hidden_size, + hidden_size, + candle_nn::Activation::Silu, + true, + "0", + "1", + )?; + let scale = 1000.0; + let half = frequency_embedding_size / 2; + let freqs = Tensor::arange(0f32, half as f32, vb.device())? + .affine(-(10000.0f64.ln()), 0.0)? + .exp()?; + Ok(Self { + mlp, + freqs, + scale, + frequency_embedding_size, + }) + } + pub fn forward(&self, t: &Tensor) -> Result { + let args = t + .affine(self.scale, 0.0)? + .unsqueeze(D::Minus1)? + .broadcast_matmul(&self.freqs.unsqueeze(0)?)?; + let mut embedding = Tensor::cat(&[args.cos()?, args.sin()?], D::Minus1)?; + if self.frequency_embedding_size % 2 > 0 { + embedding = embedding.pad_with_zeros(D::Minus1, 0, 1)?; + } + embedding = self.mlp.forward(&embedding)?; + Ok(embedding) + } +} + +pub struct SConv1d { + conv: WNConv1d, + ks: usize, + stride: usize, + dilation: usize, +} + +impl SConv1d { + pub fn new( + vb: VarBuilder, + in_c: usize, + out_c: usize, + ks: usize, + stride: usize, + dilation: usize, + groups: usize, + bias: bool, + ) -> Result { + let conv = WNConv1d::new( + vb.pp("conv.conv"), + in_c, + out_c, + ks, + dilation, + 0, + groups, + stride, + bias, + )?; + Ok(Self { + conv, + ks, + stride, + dilation, + }) + } + pub fn forward(&self, xs: &Tensor) -> Result { + let length = xs.dim(D::Minus1)?; + let ks = (self.ks - 1) * self.dilation + 1; + let padding_total = ks - self.stride; + let n_frames = (length - ks + padding_total) as f32 / self.stride as f32 + 1.0; + let idea_length = (n_frames.ceil() as usize - 1) * self.stride + (ks - padding_total); + let extra_padding = idea_length - length; + let padding_right = padding_total / 2; + let padding_left = padding_total - padding_right; + let xs = pad_reflect_last_dim(xs, (padding_left, padding_right + extra_padding))?; + let xs = self.conv.forward(&xs)?; + Ok(xs) + } +} + +pub struct Wavenet { + cond_layer: Option, + in_layers: Vec, + res_skip_layers: Vec, + hidden_c: usize, + n_layers: usize, +} + +impl Wavenet { + pub fn new( + vb: VarBuilder, + hidden_c: usize, + ks: usize, + dilation_rate: usize, + n_layers: usize, + gin_channels: usize, + ) -> Result { + let cond_layer = if gin_channels != 0 { + Some(SConv1d::new( + vb.pp("cond_layer"), + gin_channels, + 2 * hidden_c * n_layers, + 1, + 1, + 1, + 1, + true, + )?) + } else { + None + }; + let mut in_layers = vec![]; + let vb_layers = vb.pp("in_layers"); + let mut res_skip_layers = vec![]; + let vb_res_skip_layers = vb.pp("res_skip_layers"); + for i in 0..n_layers { + let dilation = dilation_rate.pow(i as u32); + let in_layer = SConv1d::new( + vb_layers.pp(i), + hidden_c, + 1 * hidden_c, + ks, + 1, + dilation, + 1, + true, + )?; + in_layers.push(in_layer); + let res_skip_c = if i < n_layers - 1 { + 2 * hidden_c + } else { + hidden_c + }; + let res_skip_layer = SConv1d::new( + vb_res_skip_layers.pp(i), + hidden_c, + res_skip_c, + 1, + 1, + 1, + 1, + true, + )?; + res_skip_layers.push(res_skip_layer); + } + Ok(Self { + cond_layer, + in_layers, + res_skip_layers, + hidden_c, + n_layers, + }) + } + + pub fn fused_add_tanh_sigmoid_multiply( + &self, + input_a: &Tensor, + input_b: &Tensor, + ) -> Result { + let in_act = input_a.add(&input_b)?; + let parts = split_tensor_with_size(&in_act, 2, 1)?; + let t_act = (&parts[0]).tanh()?; + let s_act = sigmoid(&parts[1])?; + let acts = t_act.mul(&s_act)?; + Ok(acts) + } + + pub fn forward(self, xs: &Tensor, x_mask: &Tensor, g: Option<&Tensor>) -> Result { + let mut output = xs.zeros_like()?; + let g = if let Some(g) = g + && let Some(cond_layer) = &self.cond_layer + { + Some(cond_layer.forward(g)?) + } else { + None + }; + let mut xs = xs.clone(); + for i in 0..self.n_layers { + let xs_in = &self.in_layers[i].forward(&xs)?; + let g_l = if let Some(g) = &g { + let cond_offset = i * 2 * self.hidden_c; + g.narrow(1, cond_offset, 2 * self.hidden_c)? + } else { + xs_in.zeros_like()? + }; + let acts = self.fused_add_tanh_sigmoid_multiply(&xs_in, &g_l)?; + let res_skip_act = &self.res_skip_layers[i].forward(&acts)?; + if i < self.n_layers - 1 { + let res_acts = res_skip_act.narrow(1, 0, self.hidden_c)?; + let out_acts = res_skip_act.narrow(1, self.hidden_c, self.hidden_c)?; + xs = xs.add(&res_acts)?.mul(x_mask)?; + output = output.add(&out_acts)?; + } else { + output = output.add(&res_skip_act)?; + } + } + output = output.mul(x_mask)?; + Ok(output) + } +} + +pub struct FinalLayer { + norm_final: LayerNorm, + linear: WNLinear, + ada_ln_modulation: Linear, // (silu+linear) +} + +impl FinalLayer { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + patch_size: usize, + out_c: usize, + ) -> Result { + let norm_final = get_layer_norm(vb.pp("norm_final"), 1e-6, hidden_size)?; + let linear = WNLinear::new( + vb.pp("linear"), + hidden_size, + patch_size * patch_size * out_c, + true, + )?; + let ada_ln_modulation = linear_b( + hidden_size, + 2 * hidden_size, + true, + vb.pp("adaLN_modulation.1"), + )?; + Ok(Self { + norm_final, + linear, + ada_ln_modulation, + }) + } + + pub fn forward(&self, xs: &Tensor, c: &Tensor) -> Result { + let linear_c = self.ada_ln_modulation.forward(c)?.chunk(2, 1)?; + let xs = self.norm_final.forward(xs)?; + let xs = linear_c[1] + .unsqueeze(1)? + .affine(1.0, 1.0)? + .broadcast_mul(&xs)? + .add(&linear_c[0].unsqueeze(1)?)?; + let xs = self.linear.forward(&xs)?; + Ok(xs) + } +} + +pub struct DiT { + transformer: DiTTransformer, + x_embedder: WNLinear, + cond_embedder: Embedding, + cond_projection: Linear, + t_embedder: TimestepEmbedder, + input_pos: Tensor, + t_embedder2: TimestepEmbedder, + conv1: Linear, + conv2: Conv1d, + wavenet: Wavenet, + final_layer: FinalLayer, + res_projection: Linear, + content_mask_embedder: Embedding, + skip_linear: Linear, + cond_x_merge_linear: Linear, + style_in: Option, + time_as_token: bool, + style_as_token: bool, + uvit_skip_connection: bool, +} + +impl DiT { + pub fn new(vb: VarBuilder, config: &S2MelConfig) -> Result { + let time_as_token = config.di_t.time_as_token; + let style_as_token = config.di_t.style_as_token; + let uvit_skip_connection = config.di_t.uvit_skip_connection; + let transformer_config = DiTModelArgs::new_from_dit_config(&config.di_t); + let transformer = DiTTransformer::new(vb.pp("transformer"), &transformer_config)?; + let x_embedder = WNLinear::new( + vb.pp("x_embedder"), + config.di_t.in_channels, + config.di_t.hidden_dim, + true, + )?; + let cond_embedder = embedding( + config.di_t.content_codebook_size, + config.di_t.hidden_dim, + vb.pp("cond_embedder"), + )?; + let cond_projection = linear_b( + config.di_t.content_dim, + config.di_t.hidden_dim, + true, + vb.pp("cond_projection"), + )?; + let t_embedder = TimestepEmbedder::new(vb.pp("t_embedder"), config.di_t.hidden_dim, 256)?; + let input_pos = Tensor::arange(0u32, 16384, vb.device())?; + let t_embedder2 = + TimestepEmbedder::new(vb.pp("t_embedder2"), config.wavenet.hidden_dim, 256)?; + let conv1 = linear_b( + config.di_t.hidden_dim, + config.wavenet.hidden_dim, + true, + vb.pp("conv1"), + )?; + let conv2 = get_conv1d( + vb.pp("conv2"), + config.wavenet.hidden_dim, + config.di_t.in_channels, + 1, + 0, + 1, + 1, + 1, + true, + )?; + let wavenet = Wavenet::new( + vb.pp("wavenet"), + config.wavenet.hidden_dim, + config.wavenet.kernel_size, + config.wavenet.dilation_rate, + config.wavenet.num_layers, + config.wavenet.hidden_dim, + )?; + let final_layer = FinalLayer::new( + vb.pp("final_layer"), + config.wavenet.hidden_dim, + 1, + config.wavenet.hidden_dim, + )?; + let res_projection = linear( + config.di_t.hidden_dim, + config.wavenet.hidden_dim, + vb.pp("res_projection"), + )?; + let content_mask_embedder = + embedding(1, config.di_t.hidden_dim, vb.pp("content_mask_embedder"))?; + let skip_linear = linear( + config.di_t.hidden_dim + config.di_t.in_channels, + config.di_t.hidden_dim, + vb.pp("skip_linear"), + )?; + let in_dim = if config.di_t.style_condition && !config.di_t.style_as_token { + config.di_t.hidden_dim + config.di_t.in_channels * 2 + config.style_encoder.dim + } else { + config.di_t.hidden_dim + config.di_t.in_channels * 2 + }; + let cond_x_merge_linear = + linear(in_dim, config.di_t.hidden_dim, vb.pp("cond_x_merge_linear"))?; + let style_in = if config.di_t.style_as_token { + Some(linear( + config.style_encoder.dim, + config.di_t.hidden_dim, + vb.pp("style_in"), + )?) + } else { + None + }; + Ok(Self { + transformer, + x_embedder, + cond_embedder, + cond_projection, + t_embedder, + input_pos, + t_embedder2, + conv1, + conv2, + wavenet, + final_layer, + res_projection, + content_mask_embedder, + skip_linear, + cond_x_merge_linear, + style_in, + time_as_token, + style_as_token, + uvit_skip_connection, + }) + } +} + +pub struct CFM { + estimator: DiT, +} + +pub struct MyModel { + cfm: CFM, +} + +pub struct IndexTTS2 { + cache_spk_cond: Option, + cache_s2mel_style: Option, + cache_s2mel_prompt: Option, + cache_spk_audio_prompt: Option, + cache_emo_cond: Option, + cache_emo_audio_prompt: Option, + cache_mel: Option, +} diff --git a/src/models/index_tts2/processor.rs b/src/models/index_tts2/processor.rs new file mode 100644 index 0000000..7b92580 --- /dev/null +++ b/src/models/index_tts2/processor.rs @@ -0,0 +1,199 @@ +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; +use anyhow::Result; +use candle_core::{D, DType, Device, IndexOp, Tensor, pickle::read_all_with_key}; +use candle_nn::VarBuilder; + +use crate::{ + models::{ + campplus::CAMPPlus, feature_extractor::seamless_m4t_feature_extractor::SeamlessM4TFeatureExtractor, index_tts2::config::{IndexTTS2Config, PreprocessParams}, mask_gct::model::RepCodec, w2v_bert_2_0::model::W2VBert2_0Model + }, + utils::{ + audio_utils::{ + create_hann_window, extract_audio_url, get_waveform_and_window_properties, kaldi_fbank, + kaldi_get_mel_banks, load_audio, mel_filter_bank, resample_simple, spectrogram, + torch_stft, + }, + get_vb_model_path, + tensor_utils::pad_reflect_last_dim, + }, +}; + +pub struct IndexTTS2Processor { + device: Device, + max_audio_length_seconds: usize, + feature_extractor: SeamlessM4TFeatureExtractor, + semantic_model: W2VBert2_0Model, + semantic_mean: Tensor, + semantic_std: Tensor, + semantic_codec: RepCodec, + s2mel_filters: Tensor, + s2mel_windows: Tensor, + s2mel_preprocess_params: PreprocessParams, + window_shift: usize, + window_size: usize, + padded_window_size: usize, + mel_energies: Tensor, + campplus_model: CAMPPlus, +} + +impl IndexTTS2Processor { + pub fn new( + path: &str, + save_dir: &str, + config: &IndexTTS2Config, + device: &Device, + dtype: DType, + ) -> Result { + let feature_extractor = SeamlessM4TFeatureExtractor::new( + 80, + 80, + crate::utils::tensor_utils::PaddingSide::Right, + 1.0, + 16000, + 2, + device, + )?; + let w2vbert2_path = save_dir.to_string() + "/facebook/w2v-bert-2.0"; + let semantic_model = W2VBert2_0Model::init(&w2vbert2_path, device, dtype)?; + let semantic_mean_var_path = path.to_string() + "/" + &config.w2v_stat; + let dict = read_all_with_key(semantic_mean_var_path, None)?; + let mut semantic_mean = Tensor::new(0.0, device)?.to_dtype(dtype)?; + let mut semantic_std = Tensor::new(1.0, device)?.to_dtype(dtype)?; + for (k, v) in dict { + if k.eq("mean") { + semantic_mean = v.to_device(device)?.to_dtype(dtype)?; + } else if k.eq("var") { + semantic_std = v.to_device(device)?.to_dtype(dtype)?.sqrt()?; + } + } + + let semantic_codec_path = + save_dir.to_string() + "/amphion/MaskGCT/semantic_codec/model.safetensors"; + let vb = + unsafe { VarBuilder::from_mmaped_safetensors(&[semantic_codec_path], dtype, &device)? }; + let semantic_codec = RepCodec::new(vb, &config.semantic_codec)?; + let s2mel_filters = mel_filter_bank( + config.s2mel.preprocess_params.spect_params.n_fft / 2 + 1, + config.s2mel.preprocess_params.spect_params.n_mels, + config.s2mel.preprocess_params.spect_params.fmin as f32, + config + .s2mel + .preprocess_params + .spect_params + .fmax + .unwrap_or(config.s2mel.preprocess_params.sr / 2) as f32, + config.s2mel.preprocess_params.sr as f32, + Some("slaney"), + crate::utils::audio_utils::MelScale::Slaney, + false, + device, + )? + .t()?; + let s2mel_windows = create_hann_window( + config.s2mel.preprocess_params.spect_params.win_length, + dtype, + device, + )?; + let (window_shift, window_size, padded_window_size) = + get_waveform_and_window_properties(16000, 10.0, 25.0, true)?; + let (mel_energies, _) = + kaldi_get_mel_banks(80, padded_window_size, 16000 as f32, 20.0, 0.0, device)?; + let mel_energies = mel_energies.pad_with_zeros(D::Minus1, 0, 1)?.t()?; + let campplus_model_path = save_dir.to_string() + + "/iic/speech_campplus_sv_zh-cn_16k-common/campplus_cn_common.bin"; + let campplus_vb = get_vb_model_path(campplus_model_path, dtype, device.clone(), None)?; + let campplus_model = CAMPPlus::new(campplus_vb, 80, 192, 32, 4, 128)?; + Ok(Self { + device: device.clone(), + max_audio_length_seconds: 15, + feature_extractor, + semantic_model, + semantic_mean, + semantic_std, + semantic_codec, + s2mel_filters, + s2mel_windows, + s2mel_preprocess_params: config.s2mel.preprocess_params.clone(), + window_shift, + window_size, + padded_window_size, + mel_energies, + campplus_model, + }) + } + + pub fn cut_audio(&self, audio: &Tensor, sr: usize) -> Result<(Tensor, usize)> { + let max_audio_samples = self.max_audio_length_seconds * sr; + let audio_lens = audio.dim(1)?; + let audio = if audio_lens > max_audio_samples { + audio.i((.., 0..max_audio_samples))? + } else { + audio.clone() + }; + Ok((audio, sr)) + } + + pub fn extract_audio_and_cut( + &self, + mes: &ChatCompletionParameters, + device: &Device, + ) -> Result<(Tensor, usize)> { + let audio_url_vec = extract_audio_url(mes); + let (audio, sr) = load_audio(&audio_url_vec[0], device)?; + self.cut_audio(&audio, sr) + } + + pub fn get_emb( + &self, + input_features: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let output = + self.semantic_model + .forward(input_features, attention_mask, Some(17), false)?; + let feature = &output.specify_layer_id_hidden_state.unwrap(); + let feature = feature + .broadcast_sub(&self.semantic_mean)? + .broadcast_div(&self.semantic_std)?; + Ok(feature) + } + + pub fn s2mel_spectrogram(&self, waveform: &Tensor) -> Result { + let pad = (self.s2mel_preprocess_params.spect_params.n_fft + - self.s2mel_preprocess_params.spect_params.hop_length) + / 2; + let pad_audio_22k = pad_reflect_last_dim(&waveform, (pad, pad))?; + let spec = torch_stft( + &pad_audio_22k, + self.s2mel_preprocess_params.spect_params.n_fft, + self.s2mel_preprocess_params.spect_params.hop_length, + &self.s2mel_windows, + )? + .transpose(1, 2)?; + let spec = self.s2mel_filters.broadcast_matmul(&spec)?; + let spec = spec.clamp(1e-5, f64::INFINITY)?.log()?; + Ok(spec) + } + pub fn process_info(&self, mes: &ChatCompletionParameters) -> Result<()> { + let (audio, sr) = self.extract_audio_and_cut(mes, &self.device)?; + let audio_22k = resample_simple(&audio, sr as i64, 22050)?; + let audio_16k = resample_simple(&audio, sr as i64, 16000)?; + let (audio_16k_features, audio_16k_mask) = + self.feature_extractor.call(&audio_16k, 16000, true, true)?; + let spk_cond_emb = self.get_emb(&audio_16k_features, audio_16k_mask.as_ref())?; + let (_, s_ref) = self.semantic_codec.quantize(&spk_cond_emb)?; + let ref_mel = self.s2mel_spectrogram(&audio_22k)?; + let feat = kaldi_fbank( + &audio_16k, + &self.mel_energies, + self.window_shift, + self.window_size, + self.padded_window_size, + 0.0, + )?; + let feat = feat.broadcast_sub(&feat.mean_keepdim(1)?)?; + let style = self.campplus_model.forward(&feat)?; + println!("style: {}", style); + Ok(()) + } +} diff --git a/src/models/index_tts2/utils.rs b/src/models/index_tts2/utils.rs new file mode 100644 index 0000000..e5c45fa --- /dev/null +++ b/src/models/index_tts2/utils.rs @@ -0,0 +1,18 @@ +use crate::utils::{download_model, get_default_save_dir}; + +pub async fn download_index_tts2_need_model(save_dir: Option<&str>) -> anyhow::Result<()> { + let save_dir = match save_dir { + Some(dir) => dir.to_string(), + None => get_default_save_dir().expect("Failed to get home directory"), + }; + + let w2v_bert2_0 = "facebook/w2v-bert-2.0"; + let mask_gct= "amphion/MaskGCT"; + // let campplus= "funasr/campplus"; // huggingface + let campplus = "iic/speech_campplus_sv_zh-cn_16k-common"; // modelscope + download_model(w2v_bert2_0, &save_dir, 3).await?; + download_model(mask_gct, &save_dir, 3).await?; + download_model(campplus, &save_dir, 3).await?; + + Ok(()) +} \ No newline at end of file diff --git a/src/models/mask_gct/config.rs b/src/models/mask_gct/config.rs new file mode 100644 index 0000000..a38930b --- /dev/null +++ b/src/models/mask_gct/config.rs @@ -0,0 +1,21 @@ +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct SemanticCodec { + pub codebook_size: usize, + pub hidden_size: usize, + pub codebook_dim: usize, + pub vocos_dim: usize, + pub vocos_intermediate_dim: usize, + pub vocos_num_layers: usize, + #[serde(default = "default_num_quantizers")] + pub num_quantizers: usize, + #[serde(default = "default_downsample_scale")] + pub downsample_scale: usize, +} + +fn default_num_quantizers() -> usize { + 1 +} + +fn default_downsample_scale() -> usize { + 1 +} \ No newline at end of file diff --git a/src/models/mask_gct/mod.rs b/src/models/mask_gct/mod.rs new file mode 100644 index 0000000..06008be --- /dev/null +++ b/src/models/mask_gct/mod.rs @@ -0,0 +1,2 @@ +pub mod model; +pub mod config; \ No newline at end of file diff --git a/src/models/mask_gct/model.rs b/src/models/mask_gct/model.rs new file mode 100644 index 0000000..c583e58 --- /dev/null +++ b/src/models/mask_gct/model.rs @@ -0,0 +1,344 @@ +use anyhow::Result; +use candle_core::{D, IndexOp, Tensor}; +use candle_nn::{ + Conv1d, Embedding, Init, LayerNorm, Linear, Module, VarBuilder, embedding, linear, +}; + +use crate::{ + models::{ + common::{WNConv1d, get_conv1d, get_layer_norm}, + mask_gct::config::SemanticCodec, + }, + utils::tensor_utils::{interpolate_nearest_1d, l2_normalize}, +}; + +pub struct ConvNeXtBlock { + dwconv: Conv1d, + norm: LayerNorm, + pwconv1: Linear, + pwconv2: Linear, + gamma: Option, +} + +impl ConvNeXtBlock { + pub fn new( + vb: VarBuilder, + dim: usize, + intermediate_dim: usize, + // layer_scale_init_value: f32, + ) -> Result { + let dwconv = get_conv1d(vb.pp("dwconv"), dim, dim, 7, 3, 1, 1, dim, true)?; + let norm = get_layer_norm(vb.pp("norm"), 1e-6, dim)?; + let pwconv1 = linear(dim, intermediate_dim, vb.pp("pwconv1"))?; + let pwconv2 = linear(intermediate_dim, dim, vb.pp("pwconv2"))?; + let gamma = vb.get_with_hints(dim, "gamma", Init::Const(1.0))?; + Ok(Self { + dwconv, + norm, + pwconv1, + pwconv2, + gamma: Some(gamma), + }) + } + pub fn forward(&self, xs: &Tensor) -> Result { + let residual = xs.clone(); + let xs = self.dwconv.forward(xs)?; + let xs = xs.transpose(1, 2)?; + let xs = self.norm.forward(&xs)?; + let xs = self.pwconv1.forward(&xs)?.gelu()?; + let mut xs = self.pwconv2.forward(&xs)?; + if let Some(gamma) = &self.gamma { + xs = xs.broadcast_mul(&gamma)?; + } + let xs = xs.transpose(1, 2)?; + let xs = residual.add(&xs)?; + Ok(xs) + } +} + +pub struct VocosBackbone { + embed: Conv1d, + norm: LayerNorm, + convnext: Vec, + final_layer_norm: LayerNorm, +} + +impl VocosBackbone { + pub fn new( + vb: VarBuilder, + input_channels: usize, + dim: usize, + intermediate_dim: usize, + num_layers: usize, + // layer_scale_init_value: Option, + ) -> Result { + let embed = get_conv1d(vb.pp("embed"), input_channels, dim, 7, 3, 1, 1, 1, true)?; + let norm = get_layer_norm(vb.pp("norm"), 1e-6, dim)?; + let vb_convnext = vb.pp("convnext"); + let mut convnext = vec![]; + for i in 0..num_layers { + let layer = ConvNeXtBlock::new(vb_convnext.pp(i), dim, intermediate_dim)?; + convnext.push(layer); + } + let final_layer_norm = get_layer_norm(vb.pp("final_layer_norm"), 1e-6, dim)?; + Ok(Self { + embed, + norm, + convnext, + final_layer_norm, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = self.embed.forward(xs)?; + let mut xs = self.norm.forward(&xs.transpose(1, 2)?)?.transpose(1, 2)?; + for layer in &self.convnext { + xs = layer.forward(&xs)?; + } + xs = self.final_layer_norm.forward(&xs.transpose(1, 2)?)?; + Ok(xs) + } +} + +pub struct FactorizedVectorQuantize { + use_l2_normlize: bool, + in_project: Option, + out_project: Option, + codebook: Embedding, +} + +impl FactorizedVectorQuantize { + pub fn new( + vb: VarBuilder, + input_dim: usize, + codebook_size: usize, + codebook_dim: usize, + use_l2_normlize: bool, + ) -> Result { + let (in_project, out_project) = if input_dim != codebook_dim { + let in_project = + WNConv1d::new(vb.pp("in_project"), input_dim, codebook_dim, 1, 1, 0, 1, 1, true)?; + let out_project = + WNConv1d::new(vb.pp("out_project"), codebook_dim, input_dim, 1, 1, 0, 1, 1, true)?; + (Some(in_project), Some(out_project)) + } else { + (None, None) + }; + let codebook = embedding(codebook_size, codebook_dim, vb.pp("codebook"))?; + Ok(Self { + use_l2_normlize, + in_project, + out_project, + codebook, + }) + } + + pub fn decode_latents(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> { + let (bs, len, dim) = xs.dims3()?; + let mut encodings = xs.transpose(1, 2)?.reshape((bs * dim, len))?; + let mut codebook = self.codebook.embeddings().clone(); + if self.use_l2_normlize { + encodings = l2_normalize(&encodings, 1)?; + codebook = l2_normalize(&codebook, 1)?; + } + let dist1 = encodings.powf(2.0)?.sum_keepdim(1)?; + let dist2 = encodings.affine(2.0, 0.0)?.matmul(&codebook.t()?)?; + let dist3 = codebook.powf(2.0)?.sum_keepdim(1)?.t()?; + let dist = dist1.broadcast_sub(&dist2)?.broadcast_add(&dist3)?; + let indices = dist + .affine(-1.0, 0.0)? + .argmax(1)? + .reshape((bs, ()))? + .to_dtype(candle_core::DType::U32)?; + let z_q = self.codebook.forward(&indices)?.transpose(1, 2)?; + Ok((z_q, indices)) + } + + pub fn forward(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> { + let mut xs = xs.clone(); + if let Some(in_proj) = &self.in_project { + xs = in_proj.forward(&xs)?; + } + let (z_q, indices) = self.decode_latents(&xs)?; + let mut z_q = xs.add(&z_q.sub(&xs)?)?; + if let Some(out_proj) = &self.out_project { + z_q = out_proj.forward(&z_q)?; + } + Ok((z_q, indices)) + } +} + +pub struct ResidualVQ { + num_quantizers: usize, + quantizers: Vec, +} + +impl ResidualVQ { + pub fn new( + vb: VarBuilder, + input_dim: usize, + num_quantizers: usize, + codebook_size: usize, + codebook_dim: usize, + // quantizer_type: &str, // now only surpport "fvq" + ) -> Result { + let vb_quantizers = vb.pp("quantizers"); + let mut quantizers = vec![]; + for i in 0..num_quantizers { + let quantizer = FactorizedVectorQuantize::new( + vb_quantizers.pp(i), + input_dim, + codebook_size, + codebook_dim, + true, + )?; + quantizers.push(quantizer); + } + Ok(Self { + num_quantizers, + quantizers, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + n_quantizers: Option, + ) -> Result<(Tensor, Tensor, Tensor)> { + let mut all_indices = vec![]; + let mut all_quantized = vec![]; + let n_quantizers = n_quantizers.unwrap_or(self.num_quantizers); + let mut residual = xs.clone(); + let mut quantized_out = Tensor::new(0.0f32, xs.device())?.to_dtype(xs.dtype())?; + for (i, quantizer) in (&self.quantizers).iter().enumerate() { + if i >= n_quantizers { + break; + } + let (z_q_i, indices_i) = quantizer.forward(&residual)?; + quantized_out = quantized_out.broadcast_add(&z_q_i)?; + residual = residual.sub(&z_q_i)?; + all_indices.push(indices_i); + all_quantized.push(z_q_i); + } + let all_indices = Tensor::stack(&all_indices, 0)?; + let all_quantized = Tensor::stack(&all_quantized, 0)?; + Ok((quantized_out, all_indices, all_quantized)) + } +} + +pub struct RepCodec { + downsample_scale: usize, + down: Option, + up: Option, + encoder_0: VocosBackbone, + encoder_1: Linear, + decoder_0: VocosBackbone, + decoder_1: Linear, + quantizer: ResidualVQ, +} + +impl RepCodec { + pub fn new(vb: VarBuilder, config: &SemanticCodec) -> Result { + let (down, up) = if config.downsample_scale > 1 { + let down = get_conv1d( + vb.pp("down"), + config.hidden_size, + config.hidden_size, + 3, + 1, + 2, + 1, + 1, + true, + )?; + let up = get_conv1d( + vb.pp("up"), + config.hidden_size, + config.hidden_size, + 3, + 1, + 1, + 1, + 1, + true, + )?; + (Some(down), Some(up)) + } else { + (None, None) + }; + let encoder_0 = VocosBackbone::new( + vb.pp("encoder.0"), + config.hidden_size, + config.vocos_dim, + config.vocos_intermediate_dim, + config.vocos_num_layers, + )?; + let encoder_1 = linear(config.vocos_dim, config.hidden_size, vb.pp("encoder.1"))?; + let decoder_0 = VocosBackbone::new( + vb.pp("decoder.0"), + config.hidden_size, + config.vocos_dim, + config.vocos_intermediate_dim, + config.vocos_num_layers, + )?; + let decoder_1 = linear(config.vocos_dim, config.hidden_size, vb.pp("decoder.1"))?; + let quantizer = ResidualVQ::new( + vb.pp("quantizer"), + config.hidden_size, + config.num_quantizers, + config.codebook_size, + config.codebook_dim, + )?; + Ok(Self { + downsample_scale: config.downsample_scale, + down, + up, + encoder_0, + encoder_1, + decoder_0, + decoder_1, + quantizer, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> { + let mut xs = xs.clone(); + if let Some(down) = &self.down { + xs = xs.transpose(1, 2)?; + xs = down.forward(&xs)?.gelu()?; + xs = xs.transpose(1, 2)?; + } + xs = self.encoder_0.forward(&xs.transpose(1, 2)?)?; + xs = self.encoder_1.forward(&xs)?; + xs = xs.transpose(1, 2)?; + let (quantized_out, all_indices, _) = self.quantizer.forward(&xs, None)?; + xs = self.decoder_0.forward(&quantized_out)?; + if let Some(up) = &self.up { + xs = xs.transpose(1, 2)?; + let last_dim = xs.dim(D::Minus1)?; + let target_size = last_dim * 2; + xs = interpolate_nearest_1d(&xs, target_size)?; + xs = up.forward(&xs)?.transpose(1, 2)?; + } + + Ok((xs, all_indices)) + } + + pub fn quantize(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> { + let mut xs = xs.clone(); + if let Some(down) = &self.down { + xs = xs.transpose(1, 2)?; + xs = down.forward(&xs)?.gelu()?; + xs = xs.transpose(1, 2)?; + } + xs = self.encoder_0.forward(&xs.transpose(1, 2)?)?; + xs = self.encoder_1.forward(&xs)?; + xs = xs.transpose(1, 2)?; + let (quantized_out, mut all_indices, _) = self.quantizer.forward(&xs, None)?; + if all_indices.dim(0)? == 1 { + all_indices = all_indices.squeeze(0)?; + } + let quantized_out = quantized_out.transpose(1, 2)?; + Ok((all_indices, quantized_out)) + } +} diff --git a/src/models/minicpm4/model.rs b/src/models/minicpm4/model.rs index 535f964..6c7c31f 100644 --- a/src/models/minicpm4/model.rs +++ b/src/models/minicpm4/model.rs @@ -121,6 +121,9 @@ impl MiniCPMDecoderLayer { cfg.intermediate_size, cfg.hidden_act, false, + None, + None, + None, )?; let input_layernorm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; diff --git a/src/models/mod.rs b/src/models/mod.rs index 4165669..ceedfe4 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,8 +1,12 @@ pub mod common; +pub mod campplus; pub mod deepseek_ocr; +pub mod feature_extractor; pub mod fun_asr_nano; pub mod glm_asr_nano; pub mod hunyuan_ocr; +pub mod index_tts2; +pub mod mask_gct; pub mod minicpm4; pub mod paddleocr_vl; pub mod qwen2_5vl; @@ -10,6 +14,7 @@ pub mod qwen3; pub mod qwen3vl; pub mod rmbg2_0; pub mod voxcpm; +pub mod w2v_bert_2_0; use aha_openai_dive::v1::resources::chat::{ ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, @@ -51,7 +56,7 @@ pub enum WhichModel { HunyuanOCR, #[value(name = "paddleocr-vl")] PaddleOCRVL, - #[value(name = "RMBG2.0")] + #[value(name = "rmbg2.0")] RMBG2_0, #[value(name = "voxcpm")] VoxCPM, diff --git a/src/models/qwen2_5vl/model.rs b/src/models/qwen2_5vl/model.rs index 0260532..01cf6fe 100644 --- a/src/models/qwen2_5vl/model.rs +++ b/src/models/qwen2_5vl/model.rs @@ -180,6 +180,9 @@ impl Qwen2_5VLVisionBlock { cfg.vision_config.intermediate_size, cfg.vision_config.hidden_act, true, + None, + None, + None, )?; let norm1 = rms_norm( cfg.vision_config.hidden_size, @@ -616,6 +619,9 @@ impl Qwen2_5VLTextDecoderLayer { cfg.intermediate_size, cfg.hidden_act, false, + None, + None, + None, )?; let input_layernorm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; diff --git a/src/models/qwen3/model.rs b/src/models/qwen3/model.rs index 003fd81..9543ad0 100644 --- a/src/models/qwen3/model.rs +++ b/src/models/qwen3/model.rs @@ -150,6 +150,9 @@ impl Qwen3DecoderLayer { config.intermediate_size, config.hidden_act, false, + None, + None, + None, )?; let input_layernorm = rms_norm( config.hidden_size, diff --git a/src/models/rmbg2_0/generate.rs b/src/models/rmbg2_0/generate.rs index db5db3d..1822a3d 100644 --- a/src/models/rmbg2_0/generate.rs +++ b/src/models/rmbg2_0/generate.rs @@ -49,7 +49,7 @@ impl RMBG2_0Model { img_std, device, dtype, - model_name: "RMBG2.0".to_string(), + model_name: "rmbg2.0".to_string(), }) } diff --git a/src/models/rmbg2_0/model.rs b/src/models/rmbg2_0/model.rs index edf2d1e..f47f8d0 100644 --- a/src/models/rmbg2_0/model.rs +++ b/src/models/rmbg2_0/model.rs @@ -7,7 +7,8 @@ use candle_nn::{ use crate::{ models::common::{ - TwoLinearMLP, deform_conv2d_kernel, get_batch_norm, get_conv2d, get_layer_norm, + Conv2dWithBN, TwoLinearMLP, deform_conv2d_kernel, get_batch_norm, get_conv2d, + get_layer_norm, }, utils::tensor_utils::{ get_equal_mask, index_select_2d, interpolate_bilinear, split_tensor_with_size, @@ -891,7 +892,7 @@ impl _ASPPModuleDeformable { padding, false, )?; - let bn = get_batch_norm(vb.pp("bn"), 1e-5, out_c)?; + let bn = get_batch_norm(vb.pp("bn"), 1e-5, out_c, true)?; Ok(Self { atrous_conv, bn }) } @@ -957,7 +958,8 @@ impl ASPPDeformable { 1, false, )?; - let global_avg_pool_2 = get_batch_norm(vb.pp("global_avg_pool.2"), 1e-5, in_channelster)?; + let global_avg_pool_2 = + get_batch_norm(vb.pp("global_avg_pool.2"), 1e-5, in_channelster, true)?; let conv1 = get_conv2d( vb.pp("conv1"), in_channelster * (2 + parallel_block_sizes.len()), @@ -969,7 +971,7 @@ impl ASPPDeformable { 1, false, )?; - let bn1 = get_batch_norm(vb.pp("bn1"), 1e-5, out_c)?; + let bn1 = get_batch_norm(vb.pp("bn1"), 1e-5, out_c, true)?; Ok(Self { aspp1, aspp_deforms_0, @@ -1031,8 +1033,8 @@ impl BasicDecBlk { 1, true, )?; - let bn_in = get_batch_norm(vb.pp("bn_in"), 1e-5, inter_channels)?; - let bn_out = get_batch_norm(vb.pp("bn_out"), 1e-5, out_c)?; + let bn_in = get_batch_norm(vb.pp("bn_in"), 1e-5, inter_channels, true)?; + let bn_out = get_batch_norm(vb.pp("bn_out"), 1e-5, out_c, true)?; Ok(Self { conv_in, dec_att, @@ -1072,32 +1074,6 @@ impl SimpleConvs { } } -struct Conv2dWithBN { - conv_0: Conv2d, - bn_1: BatchNorm, -} - -impl Conv2dWithBN { - pub fn new( - vb: VarBuilder, - in_c: usize, - out_c: usize, - ks: usize, - padding: usize, - stride: usize, - ) -> Result { - let conv_0 = get_conv2d(vb.pp("0"), in_c, out_c, ks, padding, stride, 1, 1, true)?; - let bn_1 = get_batch_norm(vb.pp("1"), 1e-5, out_c)?; - Ok(Self { conv_0, bn_1 }) - } - - pub fn forward(&self, x: &Tensor) -> Result { - let x = self.conv_0.forward(x)?; - let x = self.bn_1.forward_t(&x, false)?.relu()?; - Ok(x) - } -} - struct Decoder { ipt_blk5: SimpleConvs, ipt_blk4: SimpleConvs, @@ -1207,9 +1183,12 @@ impl Decoder { // let conv_ms_spvn_2 = // get_conv2d(vb.pp("conv_ms_spvn_2"), channels[3], 1, 1, 0, 1, 1, 1, true)?; let n = 16usize; - let gdt_convs_4 = Conv2dWithBN::new(vb.pp("gdt_convs_4"), channels[1], n, 3, 1, 1)?; - let gdt_convs_3 = Conv2dWithBN::new(vb.pp("gdt_convs_3"), channels[2], n, 3, 1, 1)?; - let gdt_convs_2 = Conv2dWithBN::new(vb.pp("gdt_convs_2"), channels[3], n, 3, 1, 1)?; + let gdt_convs_4 = + Conv2dWithBN::new(vb.pp("gdt_convs_4"), channels[1], n, 3, 1, 1, true, true)?; + let gdt_convs_3 = + Conv2dWithBN::new(vb.pp("gdt_convs_3"), channels[2], n, 3, 1, 1, true, true)?; + let gdt_convs_2 = + Conv2dWithBN::new(vb.pp("gdt_convs_2"), channels[3], n, 3, 1, 1, true, true)?; let gdt_convs_attn_4 = get_conv2d(vb.pp("gdt_convs_attn_4.0"), n, 1, 1, 0, 1, 1, 1, true)?; let gdt_convs_attn_3 = get_conv2d(vb.pp("gdt_convs_attn_3.0"), n, 1, 1, 0, 1, 1, 1, true)?; diff --git a/src/models/voxcpm/minicpm4.rs b/src/models/voxcpm/minicpm4.rs index 756837e..1f883b9 100644 --- a/src/models/voxcpm/minicpm4.rs +++ b/src/models/voxcpm/minicpm4.rs @@ -130,6 +130,9 @@ impl MiniCPMDecoderLayer { cfg.intermediate_size, candle_nn::Activation::Silu, false, + None, + None, + None, )?; let input_layernorm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; diff --git a/src/models/w2v_bert_2_0/config.rs b/src/models/w2v_bert_2_0/config.rs new file mode 100644 index 0000000..51142ac --- /dev/null +++ b/src/models/w2v_bert_2_0/config.rs @@ -0,0 +1,61 @@ +use candle_nn::Activation; + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct W2VBert2_0Config { + pub activation_dropout: f32, + pub adapter_act: String, + pub adapter_kernel_size: usize, + pub adapter_stride: usize, + pub add_adapter: bool, + pub apply_spec_augment: bool, + pub attention_dropout: f32, + pub bos_token_id: usize, + pub classifier_proj_size: usize, + pub codevector_dim: usize, + pub conformer_conv_dropout: f32, + pub contrastive_logits_temperature: f32, + pub conv_depthwise_kernel_size: usize, + pub ctc_loss_reduction: String, + pub ctc_zero_infinity: bool, + pub diversity_loss_weight: f32, + pub eos_token_id: usize, + pub feat_proj_dropout: f32, + pub feat_quantizer_dropout: f32, + pub feature_projection_input_dim: usize, + pub final_dropout: f32, + pub hidden_act: Activation, + pub hidden_dropout: f32, + pub hidden_size: usize, + pub initializer_range: f32, + pub intermediate_size: usize, + pub layer_norm_eps: f64, + pub layerdrop: f32, + pub left_max_position_embeddings: usize, + pub mask_feature_length: usize, + pub mask_feature_min_masks: usize, + pub mask_feature_prob: f32, + pub mask_time_length: usize, + pub mask_time_min_masks: usize, + pub mask_time_prob: f32, + pub max_source_positions: usize, + pub num_adapter_layers: usize, + pub num_attention_heads: usize, + pub num_codevector_groups: usize, + pub num_codevectors_per_group: usize, + pub num_hidden_layers: usize, + pub num_negatives: usize, + pub output_hidden_size: usize, + pub pad_token_id: usize, + pub position_embeddings_type: String, + pub proj_codevector_dim: usize, + pub right_max_position_embeddings: usize, + pub rotary_embedding_base: usize, + pub tdnn_dilation: Vec, + pub tdnn_dim: Vec, + pub tdnn_kernel: Vec, + pub torch_dtype: String, + pub use_intermediate_ffn_before_adapter: bool, + pub use_weighted_layer_sum: bool, + pub vocab_size: Option, + pub xvector_output_dim: usize, +} \ No newline at end of file diff --git a/src/models/w2v_bert_2_0/mod.rs b/src/models/w2v_bert_2_0/mod.rs new file mode 100644 index 0000000..3621b7a --- /dev/null +++ b/src/models/w2v_bert_2_0/mod.rs @@ -0,0 +1,2 @@ +pub mod config; +pub mod model; \ No newline at end of file diff --git a/src/models/w2v_bert_2_0/model.rs b/src/models/w2v_bert_2_0/model.rs new file mode 100644 index 0000000..6b5b6e5 --- /dev/null +++ b/src/models/w2v_bert_2_0/model.rs @@ -0,0 +1,571 @@ +use anyhow::Result; +use candle_core::{D, DType, Device, Tensor}; +use candle_nn::{ + Activation, Conv1d, Embedding, Init, LayerNorm, Linear, Module, VarBuilder, embedding, linear, + linear_b, +}; + +use crate::{ + models::{ + common::{ + GLU, NaiveAttention, TwoLinearMLP, eager_attention_forward, get_conv1d, get_layer_norm, + }, + w2v_bert_2_0::config::W2VBert2_0Config, + }, + position_embed::rope::{RoPE, apply_rotary_pos_emb}, + utils::{find_type_files, tensor_utils::masked_fill_zeros}, +}; + +pub struct Wav2Vec2BertFeatureProjection { + layer_norm: LayerNorm, + projection: Linear, +} + +impl Wav2Vec2BertFeatureProjection { + pub fn new(vb: VarBuilder, config: &W2VBert2_0Config) -> Result { + let layer_norm = get_layer_norm( + vb.pp("layer_norm"), + config.layer_norm_eps, + config.feature_projection_input_dim, + )?; + let projection = linear( + config.feature_projection_input_dim, + config.hidden_size, + vb.pp("projection"), + )?; + Ok(Self { + layer_norm, + projection, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> { + let norm_xs = self.layer_norm.forward(xs)?; + let xs = self.projection.forward(&norm_xs)?; + Ok((xs, norm_xs)) + } +} + +pub struct Wav2Vec2BertSelfAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + head_dim: usize, + num_heads: usize, + position_embeddings_type: Option, + linear_pos: Option, + pos_bias_u: Option, + pos_bias_v: Option, + left_max_position_embeddings: usize, + right_max_position_embeddings: usize, + distance_embedding: Option, +} + +impl Wav2Vec2BertSelfAttention { + pub fn new( + vb: VarBuilder, + config: &W2VBert2_0Config, + is_adapter_attention: bool, + ) -> Result { + let hidden_size = if is_adapter_attention { + config.hidden_size + } else { + config.output_hidden_size + }; + let head_dim = hidden_size / config.num_attention_heads; + let num_heads = config.num_attention_heads; + let left_max_position_embeddings = config.left_max_position_embeddings; + let right_max_position_embeddings = config.right_max_position_embeddings; + let position_embeddings_type = if !is_adapter_attention { + Some(config.position_embeddings_type.clone()) + } else { + None + }; + let (linear_pos, pos_bias_u, pos_bias_v, distance_embedding) = + if let Some(pos_type) = &position_embeddings_type { + if pos_type.eq("relative") { + let linear_pos = Some(linear_b( + hidden_size, + hidden_size, + false, + vb.pp("linear_pos"), + )?); + let pos_bias_u = Some(vb.get_with_hints( + (config.num_attention_heads, head_dim), + "pos_bias_u", + Init::Const(0.), + )?); + let pos_bias_v = Some(vb.get_with_hints( + (config.num_attention_heads, head_dim), + "pos_bias_v", + Init::Const(0.), + )?); + (linear_pos, pos_bias_u, pos_bias_v, None) + } else if pos_type.eq("relative_key") { + let num_positions = + left_max_position_embeddings + right_max_position_embeddings + 1; + let distance_embedding = Some(embedding( + num_positions, + head_dim, + vb.pp("distance_embedding"), + )?); + (None, None, None, distance_embedding) + } else { + (None, None, None, None) + } + } else { + (None, None, None, None) + }; + let q_proj = linear_b(hidden_size, hidden_size, true, vb.pp("linear_q"))?; + let k_proj = linear_b(hidden_size, hidden_size, true, vb.pp("linear_k"))?; + let v_proj = linear_b(hidden_size, hidden_size, true, vb.pp("linear_v"))?; + let o_proj = linear_b(hidden_size, hidden_size, true, vb.pp("linear_out"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + head_dim, + num_heads, + position_embeddings_type, + linear_pos, + pos_bias_u, + pos_bias_v, + left_max_position_embeddings, + right_max_position_embeddings, + distance_embedding, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + cos: Option<&Tensor>, + sin: Option<&Tensor>, + attention_mask: Option<&Tensor>, + ) -> Result { + if let Some(pos_type) = &self.position_embeddings_type + && pos_type.eq("rotary") + && (cos.is_none() || sin.is_none()) + { + return Err(anyhow::anyhow!( + "rotary type position cos and sin can not be none" + )); + } + + let (b_sz, q_len, _) = xs.dims3()?; + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + let query_states = query_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let (query_states, key_states) = if let Some(cos) = cos + && let Some(sin) = sin + { + apply_rotary_pos_emb(&query_states, &key_states, cos, sin, false)? + } else { + (query_states, key_states) + }; + let scale = 1f64 / f64::sqrt(self.head_dim as f64); + let attention_mask = if let Some(pos_type) = &self.position_embeddings_type + && pos_type.eq("relative_key") + && let Some(embed) = &self.distance_embedding + { + let query_length = query_states.dim(2)?; + let key_length = key_states.dim(2)?; + let position_ids_l = + Tensor::arange(0i64, query_length as i64, xs.device())?.unsqueeze(D::Minus1)?; + let position_ids_r = + Tensor::arange(0i64, key_length as i64, xs.device())?.unsqueeze(0)?; + let distance = position_ids_r.broadcast_sub(&position_ids_l)?; + let distance = distance.clamp( + -(self.left_max_position_embeddings as i64), + self.right_max_position_embeddings as i64, + )?; + let distance = distance + .affine(1.0, self.left_max_position_embeddings as f64)? + .to_dtype(candle_core::DType::U32)?; + let pos_emb = embed.forward(&distance)?.to_dtype(query_states.dtype())?; // (seq_q, seq_k, dim) + let query_ = query_states.unsqueeze(D::Minus2)?; // (b, n_head, seq_q, 1, dim) + let pos_emb = pos_emb.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, se_q, seq_k, dim) + // torch.einsum("bhld,lrd->bhlr", query, positional_embedding) + // (bs, n_head, seq_len, seq_len) + let relative_position_attn_weights = query_ + .broadcast_mul(&pos_emb)? + .sum(D::Minus1)? + .affine(scale, 0.0)?; + if let Some(mask) = attention_mask { + // let mask = mask.unsqueeze(1)?.unsqueeze(D::Minus1)?; + Some(relative_position_attn_weights.broadcast_add(&mask)?) + } else { + Some(relative_position_attn_weights) + } + } else { + if let Some(mask) = attention_mask { + // let mask = mask.unsqueeze(1)?.unsqueeze(D::Minus1)?; + Some(mask.clone()) + } else { + None + } + }; + + let attn_output = eager_attention_forward( + &query_states, + &key_states, + &value_states, + None, + attention_mask.as_ref(), + scale, + )?; + let attn_output = attn_output.reshape((b_sz, q_len, self.num_heads * self.head_dim))?; + let attn_output = attn_output.apply(&self.o_proj)?; + Ok(attn_output) + } +} + +pub struct Wav2Vec2BertConvolutionModule { + layer_norm: LayerNorm, + pointwise_conv1: Conv1d, + glu: GLU, + conv_depthwise_kernel_size: usize, + depthwise_conv: Conv1d, + depthwise_layer_norm: LayerNorm, + act: Activation, + pointwise_conv2: Conv1d, +} + +impl Wav2Vec2BertConvolutionModule { + pub fn new(vb: VarBuilder, config: &W2VBert2_0Config) -> Result { + let layer_norm = get_layer_norm( + vb.pp("layer_norm"), + config.layer_norm_eps, + config.hidden_size, + )?; + let pointwise_conv1 = get_conv1d( + vb.pp("pointwise_conv1"), + config.hidden_size, + 2 * config.hidden_size, + 1, + 0, + 1, + 1, + 1, + false, + )?; + let glu = GLU::new(1)?; + let conv_depthwise_kernel_size = config.conv_depthwise_kernel_size; + let depthwise_conv = get_conv1d( + vb.pp("depthwise_conv"), + config.hidden_size, + config.hidden_size, + conv_depthwise_kernel_size, + 0, + 1, + 1, + config.hidden_size, + false, + )?; + let depthwise_layer_norm = get_layer_norm( + vb.pp("depthwise_layer_norm"), + config.layer_norm_eps, + config.hidden_size, + )?; + let pointwise_conv2 = get_conv1d( + vb.pp("pointwise_conv2"), + config.hidden_size, + config.hidden_size, + 1, + 0, + 1, + 1, + 1, + false, + )?; + Ok(Self { + layer_norm, + pointwise_conv1, + glu, + conv_depthwise_kernel_size, + depthwise_conv, + depthwise_layer_norm, + act: config.hidden_act, + pointwise_conv2, + }) + } + + pub fn forward(&self, xs: &Tensor, mask: Option<&Tensor>) -> Result { + let mut xs = self.layer_norm.forward(xs)?; + if let Some(mask) = mask { + xs = masked_fill_zeros(&xs, mask)?; + } + let xs = xs.transpose(1, 2)?; + // (batch, 2*channel, dim) + let xs = self.pointwise_conv1.forward(&xs)?; + // (batch, channel, dim) + let xs = self.glu.forward(&xs)?; + let xs = xs.pad_with_zeros(D::Minus1, self.conv_depthwise_kernel_size - 1, 0)?; + let xs = self.depthwise_conv.forward(&xs)?; + let xs = self + .depthwise_layer_norm + .forward(&xs.transpose(1, 2)?)? + .transpose(1, 2)?; + let xs = xs.apply(&self.act)?; + let xs = self.pointwise_conv2.forward(&xs)?; + let xs = xs.transpose(1, 2)?; + Ok(xs) + } +} + +pub struct Wav2Vec2BertEncoderLayer { + ffn1_layer_norm: LayerNorm, + ffn1: TwoLinearMLP, + self_attn_layer_norm: LayerNorm, + self_attn: Wav2Vec2BertSelfAttention, + conv_module: Wav2Vec2BertConvolutionModule, + ffn2_layer_norm: LayerNorm, + ffn2: TwoLinearMLP, + final_layer_norm: LayerNorm, +} + +impl Wav2Vec2BertEncoderLayer { + pub fn new(vb: VarBuilder, config: &W2VBert2_0Config) -> Result { + let ffn1_layer_norm = get_layer_norm( + vb.pp("ffn1_layer_norm"), + config.layer_norm_eps, + config.hidden_size, + )?; + let ffn1 = TwoLinearMLP::new( + vb.pp("ffn1"), + config.hidden_size, + config.intermediate_size, + config.hidden_size, + config.hidden_act, + true, + "intermediate_dense", + "output_dense", + )?; + let self_attn_layer_norm = get_layer_norm( + vb.pp("self_attn_layer_norm"), + config.layer_norm_eps, + config.hidden_size, + )?; + let self_attn = Wav2Vec2BertSelfAttention::new(vb.pp("self_attn"), config, false)?; + let conv_module = Wav2Vec2BertConvolutionModule::new(vb.pp("conv_module"), config)?; + let ffn2_layer_norm = get_layer_norm( + vb.pp("ffn2_layer_norm"), + config.layer_norm_eps, + config.hidden_size, + )?; + let ffn2 = TwoLinearMLP::new( + vb.pp("ffn2"), + config.hidden_size, + config.intermediate_size, + config.hidden_size, + config.hidden_act, + true, + "intermediate_dense", + "output_dense", + )?; + let final_layer_norm = get_layer_norm( + vb.pp("final_layer_norm"), + config.layer_norm_eps, + config.hidden_size, + )?; + Ok(Self { + ffn1_layer_norm, + ffn1, + self_attn_layer_norm, + self_attn, + conv_module, + ffn2_layer_norm, + ffn2, + final_layer_norm, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + cos: Option<&Tensor>, + sin: Option<&Tensor>, + attention_mask: Option<&Tensor>, + conv_attention_mask: Option<&Tensor>, + ) -> Result { + let residual = xs.clone(); + let xs = self.ffn1_layer_norm.forward(xs)?; + let xs = self.ffn1.forward(&xs)?; + let residual = xs.affine(0.5, 0.0)?.add(&residual)?; + let xs = self.self_attn_layer_norm.forward(&residual)?; + let xs = self.self_attn.forward(&xs, cos, sin, attention_mask)?; + let residual = xs.add(&residual)?; + let xs = self.conv_module.forward(&residual, conv_attention_mask)?; + let residual = xs.add(&residual)?; + let xs = self.ffn2_layer_norm.forward(&residual)?; + let xs = self.ffn2.forward(&xs)?; + let xs = xs.affine(0.5, 0.0)?.add(&residual)?; + let xs = self.final_layer_norm.forward(&xs)?; + Ok(xs) + } +} + +pub struct ModelOutput { + pub last_hidden_state: Tensor, + pub specify_layer_id_hidden_state: Option, + pub hidden_states: Option>, +} + +pub struct Wav2Vec2BertEncoder { + embed_positions: Option, + layers: Vec, +} + +impl Wav2Vec2BertEncoder { + pub fn new(vb: VarBuilder, config: &W2VBert2_0Config) -> Result { + let embed_positions = if config.position_embeddings_type.eq("rotary") { + let dim = config.hidden_size / config.num_attention_heads; + let embed_positions = RoPE::new(dim, 10000.0, vb.device())?; + Some(embed_positions) + } else { + None + }; + let vb_layers = vb.pp("layers"); + let mut layers = vec![]; + for i in 0..config.num_hidden_layers { + let layer = Wav2Vec2BertEncoderLayer::new(vb_layers.pp(i), config)?; + layers.push(layer); + } + Ok(Self { + embed_positions, + layers, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + layer_id: Option, + output_hidden_states: bool, + ) -> Result { + // xs: (bs, seq_len ,dim) + // attention_mask: Some: (bs, seq_len) + let (_, seq_len, _) = xs.dims3()?; + let conv_attention_mask = attention_mask; + let (mut xs, attention_mask) = if let Some(mask) = attention_mask { + let xs = masked_fill_zeros(xs, mask)?; + // (bs, 1, 1, seq_len) + let attention_mask = mask.unsqueeze(1)?.unsqueeze(1)?; + let neg_inf_t = attention_mask + .zeros_like()? + .to_dtype(xs.dtype())? + .affine(1.0, f64::NEG_INFINITY)?; + let attention_mask_f = attention_mask.to_dtype(xs.dtype())?; + let attention_mask = attention_mask + .where_cond(&attention_mask_f, &neg_inf_t)? + .to_dtype(xs.dtype())? + .affine(1.0, -1.0)?; + (xs, Some(attention_mask)) + } else { + (xs.clone(), None) + }; + let (cos, sin) = if let Some(embed_posi) = &self.embed_positions { + let (cos, sin) = embed_posi.forward(0, seq_len, xs.device())?; + (Some(cos), Some(sin)) + } else { + (None, None) + }; + let mut hidden_states: Vec = vec![]; + let mut specify_layer_id_hidden_state = None; + + for (i, layer) in (&self.layers).iter().enumerate() { + if output_hidden_states { + hidden_states.push(xs.clone()); + } + if let Some(id) = layer_id + && id == i + { + specify_layer_id_hidden_state = Some(xs.clone()); + } + xs = layer.forward( + &xs, + cos.as_ref(), + sin.as_ref(), + attention_mask.as_ref(), + conv_attention_mask, + )?; + } + let hidden_states = if hidden_states.len() > 0 { + Some(hidden_states) + } else { + None + }; + Ok(ModelOutput { + last_hidden_state: xs, + specify_layer_id_hidden_state, + hidden_states, + }) + } +} + +pub struct W2VBert2_0Model { + config: W2VBert2_0Config, + feature_projection: Wav2Vec2BertFeatureProjection, + masked_spec_embed: Option, + encoder: Wav2Vec2BertEncoder, + // config.add_adapter is false, adapter is None, Wav2Vec2BertAdapter not complish + // adapter: Option, + // config.use_intermediate_ffn_before_adapter is false, intermediate_ffn is None + // intermediate_ffn: Option, +} + +impl W2VBert2_0Model { + pub fn init(path: &str, device: &Device, dtype: DType) -> Result { + let config_path = path.to_string() + "/config.json"; + let config: W2VBert2_0Config = serde_json::from_slice(&std::fs::read(config_path)?)?; + let model_list = find_type_files(path, "safetensors")?; + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? }; + W2VBert2_0Model::new(vb, &config) + } + + pub fn new(vb: VarBuilder, config: &W2VBert2_0Config) -> Result { + let feature_projection = + Wav2Vec2BertFeatureProjection::new(vb.pp("feature_projection"), config)?; + let masked_spec_embed = if config.mask_time_prob > 0.0 || config.mask_time_prob > 0.0 { + Some( + vb.get_with_hints(config.hidden_size, "masked_spec_embed", Init::Uniform { + lo: 0.0, + up: 1.0, + })?, + ) + } else { + None + }; + let encoder = Wav2Vec2BertEncoder::new(vb.pp("encoder"), config)?; + Ok(Self { + config: config.clone(), + feature_projection, + masked_spec_embed, + encoder, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + attention_mask: Option<&Tensor>, + layer_id: Option, + output_hidden_states: bool, + ) -> Result { + let (xs, _) = self.feature_projection.forward(xs)?; + self.encoder + .forward(&xs, attention_mask, layer_id, output_hidden_states) + } +} diff --git a/src/position_embed/rope.rs b/src/position_embed/rope.rs index 8853489..7979c6c 100644 --- a/src/position_embed/rope.rs +++ b/src/position_embed/rope.rs @@ -1,4 +1,4 @@ -use anyhow::Result; +use anyhow::{Result, anyhow}; use candle_core::{D, DType, Device, IndexOp, Tensor}; use candle_transformers::models::deepseek2::SplitOp; @@ -158,6 +158,68 @@ pub fn glm_asr_apply_rotary_pos_emb( Ok((q_embed, k_embed)) } +pub fn roformer_rotate(x: &Tensor) -> Result { + let dims = x.dims(); + let last_dim = dims + .last() + .ok_or(anyhow!("Input tensor must have at least one dimension"))?; + if last_dim % 2 != 0 { + return Err(anyhow!( + "Last dimension size must be even, got {}", + last_dim + )); + } + let new_dims: Vec = dims[..dims.len() - 1] + .iter() + .copied() + .chain([last_dim / 2, 2]) + .collect(); + let x_reshape = x.reshape(new_dims)?; + let x_chunks = x_reshape.chunk(2, D::Minus1)?; + let x1 = &x_chunks[0]; + let x2 = &x_chunks[1]; + // let x1 = x_reshape.narrow(D::Minus1, 0, 1)?; + // let x2 = x_reshape.narrow(D::Minus1, 1, 1)?; + let x2_neg = x2.affine(-1.0, 0.0)?; + let rotate_x = Tensor::cat(&[&x2_neg, x1], D::Minus1)?; + Ok(rotate_x.flatten(D::Minus2, D::Minus1)?) +} + +pub fn apply_rotary_pos_emb_roformer( + q: &Tensor, + k: &Tensor, + cos: &Tensor, + sin: &Tensor, + tof32: bool, +) -> Result<(Tensor, Tensor)> { + let mut cos = cos.clone(); + let mut sin = sin.clone(); + if cos.rank() == 2 { + // (seq_len, head_dim) -> (1, 1, seq_len, head_dim) + cos = cos.unsqueeze(0)?.unsqueeze(0)?; + sin = sin.unsqueeze(0)?.unsqueeze(0)?; + } + if cos.rank() == 3 { + // (bs, seq_len, head_dim) -> (bs, 1, seq_len, head_dim) + cos = cos.unsqueeze(1)?; + sin = sin.unsqueeze(1)?; + } + let orig_dtype = q.dtype(); + let q = if tof32 { &q.to_dtype(DType::F32)? } else { q }; + let k = if tof32 { &k.to_dtype(DType::F32)? } else { k }; + let cos = cos.to_dtype(q.dtype())?; + let sin = sin.to_dtype(q.dtype())?; + let q_embed = q + .broadcast_mul(&cos)? + .add(&roformer_rotate(q)?.broadcast_mul(&sin)?)? + .to_dtype(orig_dtype)?; + let k_embed = k + .broadcast_mul(&cos)? + .add(&roformer_rotate(k)?.broadcast_mul(&sin)?)? + .to_dtype(orig_dtype)?; + Ok((q_embed, k_embed)) +} + #[derive(Debug, Clone)] pub struct Qwen2_5VLTextRotaryEmbedding { inv_freq: Vec, diff --git a/src/utils/audio_utils.rs b/src/utils/audio_utils.rs index 0a569b3..3b89a72 100644 --- a/src/utils/audio_utils.rs +++ b/src/utils/audio_utils.rs @@ -32,7 +32,7 @@ use symphonia::core::meta::MetadataOptions; use symphonia::core::probe::Hint; use crate::utils::get_default_save_dir; -use crate::utils::tensor_utils::{linspace, pad_replicate_last_dim}; +use crate::utils::tensor_utils::{linspace, log10, pad_reflect_last_dim, pad_replicate_last_dim}; // 重采样方法枚举 #[derive(Debug, Clone, Copy)] @@ -569,6 +569,11 @@ pub fn load_audio_use_symphonia(audio_vec: Vec, device: &Device) -> Result<( Ok((audio_tensor, sample_rate as usize)) } +pub fn load_audio(path: &str, device: &Device) -> Result<(Tensor, usize)> { + let audio_vec = get_audio_bytes_vec(path)?; + load_audio_use_symphonia(audio_vec, device) +} + pub fn load_audio_with_resample( path: &str, device: &Device, @@ -945,18 +950,44 @@ pub fn load_and_resample_audio_ffmpeg( // Ok(audio) // } +// pub fn create_hann_window(window_size: usize, dtype: DType, device: &Device) -> Result { +// let n = window_size as f64; +// let window: Vec = (0..window_size) +// .map(|i| { +// let i_f64 = i as f64; +// let val = 0.5 * (1.0 - (2.0 * PI * i_f64 / n).cos()); +// val as f32 +// }) +// .collect(); +// Ok(Tensor::from_vec(window, window_size, device)?.to_dtype(dtype)?) +// } + pub fn create_hann_window(window_size: usize, dtype: DType, device: &Device) -> Result { - let n = window_size as f64; - let window: Vec = (0..window_size) + if window_size < 1 { + return Err(anyhow::anyhow!("window_size must bigger than 0")); + } + if window_size == 1 { + return Ok(Tensor::new(1.0f32, device)?.to_dtype(dtype)?); + } + let n = window_size as f64 - 1.0; + let start = 1_i64 - window_size as i64; + let end = window_size as i64; + let window: Vec = (start..end) + .step_by(2) .map(|i| { let i_f64 = i as f64; - let val = 0.5 * (1.0 - (2.0 * PI * i_f64 / n).cos()); + let val = 0.5 + 0.5 * (PI * i_f64 / n).cos(); val as f32 }) .collect(); Ok(Tensor::from_vec(window, window_size, device)?.to_dtype(dtype)?) } +pub fn create_povey_window(window_size: usize, dtype: DType, device: &Device) -> Result { + let window = create_hann_window(window_size, dtype, device)?; + Ok(window.powf(0.85)?) +} + pub fn crate_hamming_window( window_size: usize, periodic: bool, @@ -1166,6 +1197,22 @@ pub fn apply_stft(waveform: &Tensor) -> Result { Ok(magnitudes) } +pub fn torch_stft( + waveform: &Tensor, + n_fft: usize, + hop_length: usize, + window: &Tensor, +) -> Result { + // waveform: already padding + // (bs, n_frames, n_fft) + let frames = extract_frames(&waveform, n_fft, hop_length)?; + // 应用汉明窗口 + let result = frames.broadcast_mul(window)?; + // 傅立叶变换 + let magnitudes = apply_stft(&result)?; + Ok(magnitudes) +} + pub fn kaldi_fbank( waveform: &Tensor, mel_energies: &Tensor, @@ -1489,3 +1536,73 @@ pub fn kaldi_get_mel_banks( Ok((bins_tensor, center_freqs)) } + +pub fn spectrogram( + waveform: &Tensor, + window: &Tensor, + frame_length: usize, + hop_length: usize, + fft_length: usize, + power: Option, + center: bool, + preemphasis: f64, + mel_filters: Option<&Tensor>, + log_mel: Option<&str>, + mel_floor: f32, + remove_dc_offset: bool, +) -> Result { + let waveform = if center { + let pad = frame_length / 2; + pad_reflect_last_dim(waveform, (pad, pad))? + } else { + waveform.clone() + }; + let mut frames = extract_frames(&waveform, frame_length, hop_length)?; + if remove_dc_offset { + let row_means = frames.mean_keepdim(D::Minus1)?; + frames = frames.broadcast_sub(&row_means)?; + } + + if preemphasis != 0.0 { + let buffer_0 = frames + .i((.., .., 0))? + .affine(1.0 - preemphasis, 0.0)? + .unsqueeze(D::Minus1)?; + let buffer_ = frames.i((.., .., 1..))?.sub( + &frames + .i((.., .., 0..frame_length - 1))? + .affine(preemphasis, 0.0)?, + )?; + frames = Tensor::cat(&[buffer_0, buffer_], D::Minus1)?; + } + let mut frames = frames.broadcast_mul(&window)?; + let pad_len = fft_length - frame_length; + if pad_len > 0 { + // (bs, nframes, frame_length) -> (bs, nframes, fft_length) + frames = frames.pad_with_zeros(D::Minus1, 0, pad_len)?; + } + let mut spectrogram = apply_stft(&frames)?; // stft已经做了pow(2.0) + spectrogram = spectrogram.transpose(D::Minus1, D::Minus2)?; + if let Some(mel_filters) = mel_filters { + let spect = mel_filters.t()?.broadcast_matmul(&spectrogram)?; + spectrogram = spect.maximum( + &Tensor::new(mel_floor, spect.device())? + .to_dtype(spect.dtype())? + .broadcast_as(spect.shape())?, + )?; + } + if let Some(_) = power + && let Some(log_mel) = log_mel + { + if log_mel == "log" { + spectrogram = spectrogram.log()?; + } else if log_mel == "log10" { + spectrogram = log10(&spectrogram)?; + } else { + return Err(anyhow!( + "dB not completed or Unknown log_mel option ".to_string() + )); + } + } + Ok(spectrogram) +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 32995f0..d8155c5 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -3,7 +3,8 @@ pub mod img_utils; pub mod tensor_utils; pub mod video_utils; -use std::{fs, process::Command}; +use std::io::Read; +use std::{collections::HashMap, fs, process::Command, time::Duration}; use aha_openai_dive::v1::resources::{ chat::{ @@ -14,10 +15,18 @@ use aha_openai_dive::v1::resources::{ }, shared::{FinishReason, Usage}, }; -use anyhow::Result; -use candle_core::{DType, Device}; +use anyhow::{Result, anyhow}; +use byteorder::{LittleEndian, ReadBytesExt}; +use candle_core::{ + Context, DType, Device, Shape, Tensor, + pickle::{Object, PthTensors, Stack, TensorInfo, read_all_with_key}, +}; +use candle_nn::VarBuilder; use candle_transformers::generation::{LogitsProcessor, Sampling}; use dirs::home_dir; +use half::{bf16, f16, slice::HalfFloatSliceExt}; +use modelscope::ModelScope; +use tokio::time::sleep; pub fn get_device(device: Option<&Device>) -> Device { match device { @@ -128,6 +137,228 @@ pub fn find_type_files(path: &str, extension_type: &str) -> Result> Ok(files) } +pub fn get_vb_model_path( + model_path: String, + dtype: DType, + device: Device, + key: Option<&'_ str>, +) -> Result> { + let mut dict_to_hashmap = HashMap::new(); + let dict = read_all_with_key(&model_path, key)?; + for (k, v) in dict { + dict_to_hashmap.insert(k, v); + } + let vb = VarBuilder::from_tensors(dict_to_hashmap, dtype, &device); + Ok(vb) +} + +pub fn get_vb_extension( + path: String, + extension_type: String, + dtype: DType, + device: Device, + key: Option<&'_ str>, +) -> Result> { + let model_list = find_type_files(&path, &extension_type)?; + let mut dict_to_hashmap = HashMap::new(); + for m in model_list { + let dict = read_all_with_key(m, key)?; + for (k, v) in dict { + dict_to_hashmap.insert(k, v); + } + } + let vb = VarBuilder::from_tensors(dict_to_hashmap, dtype, &device); + Ok(vb) +} + +pub fn crate_tensor_from_reader( + shape: Shape, + dtype: DType, + reader: &mut R, +) -> Result { + let elem_count = shape.elem_count(); + match dtype { + DType::BF16 => { + let mut data_t = vec![bf16::ZERO; elem_count]; + reader.read_u16_into::(data_t.reinterpret_cast_mut())?; + Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?) + } + DType::F16 => { + let mut data_t = vec![f16::ZERO; elem_count]; + reader.read_u16_into::(data_t.reinterpret_cast_mut())?; + Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?) + } + DType::F32 => { + let mut data_t = vec![0f32; elem_count]; + reader.read_f32_into::(&mut data_t)?; + Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?) + } + DType::F64 => { + let mut data_t = vec![0f64; elem_count]; + reader.read_f64_into::(&mut data_t)?; + Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?) + } + DType::U8 => { + let mut data_t = vec![0u8; elem_count]; + reader.read_exact(&mut data_t)?; + Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?) + } + DType::U32 => { + let mut data_t = vec![0u32; elem_count]; + reader.read_u32_into::(&mut data_t)?; + Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?) + } + DType::I64 => { + let mut data_t = vec![0i64; elem_count]; + reader.read_i64_into::(&mut data_t)?; + Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?) + } + } +} + +pub fn read_pth_tensor_info_cycle>( + path: P, + key: Option<&str>, +) -> Result> { + let file = std::fs::File::open(path.as_ref())?; + let zip_reader = std::io::BufReader::new(file); + let mut zip = zip::ZipArchive::new(zip_reader)?; + let zip_file_names = zip + .file_names() + .map(|f| f.to_string()) + .collect::>(); + + let mut tensor_infos = vec![]; + for file_name in zip_file_names.iter() { + if !file_name.ends_with("data.pkl") { + continue; + } + let dir_name = std::path::PathBuf::from(file_name.strip_suffix(".pkl").context("no .pkl")?); + let reader = zip.by_name(file_name)?; + let mut reader = std::io::BufReader::new(reader); + let mut stack = Stack::empty(); + stack.read_loop(&mut reader)?; + let obj = stack.finalize()?; + + let obj = match obj { + Object::Build { callable, args } => match *callable { + Object::Reduce { callable, args: _ } => match *callable { + Object::Class { + module_name, + class_name, + } if module_name == "__torch__" && class_name == "Module" => *args, + _ => continue, + }, + _ => continue, + }, + obj => obj, + }; + + // If key is provided, then we need to extract the state_dict from the object. + let obj = if let Some(key) = key { + let multi_key: Vec<&str> = key.split(".").collect(); + if multi_key.len() > 1 { + let mut current_obj = obj; + for k in multi_key.iter() { + if let Object::Dict(key_values) = current_obj { + current_obj = key_values + .into_iter() + .find(|(key_obj, _)| *key_obj == Object::Unicode(k.to_string())) + .map(|(_, v)| v) + .ok_or_else(|| anyhow!(format!("key '{}' not found", k)))?; + } else { + return Err(anyhow!(format!( + "Expected dictionary at key '{}', but found other type", + k + ))); + } + } + current_obj + } else { + if let Object::Dict(key_values) = obj { + key_values + .into_iter() + .find(|(k, _)| *k == Object::Unicode(key.to_owned())) + .map(|(_, v)| v) + .ok_or_else(|| anyhow!(format!("key {key} not found")))? + } else { + obj + } + } + } else { + obj + }; + + // If the object is a dict, then we can extract the tensor info from it. + // NOTE: We are assuming that the `obj` is state_dict by this stage. + if let Object::Dict(key_values) = obj { + for (name, value) in key_values.into_iter() { + match value.into_tensor_info(name, &dir_name) { + Ok(Some(tensor_info)) => tensor_infos.push(tensor_info), + Ok(None) => {} + Err(err) => eprintln!("skipping: {err:?}"), + } + } + } + } + let tensor_infos: HashMap = tensor_infos + .into_iter() + .map(|ti| (ti.name.to_string(), ti)) + .collect(); + + let tensor_names = tensor_infos.keys(); + let mut tensors = Vec::with_capacity(tensor_names.len()); + for name in tensor_names { + let _ = match tensor_infos.get(name) { + None => {} + Some(tensor_info) => { + let zip_reader = std::io::BufReader::new(std::fs::File::open(&path)?); + let mut zip = zip::ZipArchive::new(zip_reader)?; + let mut reader = zip.by_name(&tensor_info.path)?; + let is_fortran_contiguous = tensor_info.layout.is_fortran_contiguous(); + let rank = tensor_info.layout.shape().rank(); + + // Reading the data is a bit tricky as it can be strided, for now only support the basic + // case and when the tensor is fortran contiguous. + if !tensor_info.layout.is_contiguous() && !is_fortran_contiguous { + return Err(anyhow!(format!( + "cannot retrieve non-contiguous tensors {:?}", + tensor_info.layout + ))); + } + let start_offset = tensor_info.layout.start_offset(); + if start_offset > 0 { + std::io::copy( + &mut reader.by_ref().take(start_offset as u64), + &mut std::io::sink(), + )?; + } + let tensor = crate_tensor_from_reader( + tensor_info.layout.shape().clone(), + tensor_info.dtype, + &mut reader, + )?; + + if rank > 1 && is_fortran_contiguous { + // Reverse the shape, e.g. Shape(2, 3, 4) -> Shape(4, 3, 2) + let shape_reversed: Vec<_> = + tensor_info.layout.dims().iter().rev().cloned().collect(); + let tensor = tensor.reshape(shape_reversed)?; + + // Permute (transpose) the dimensions, e.g. Shape(4, 3, 2) -> Shape(2, 3, 4) + let dim_indeces_reversed: Vec<_> = (0..rank).rev().collect(); + let tensor = tensor.permute(dim_indeces_reversed)?; + // Ok(Some(tensor)) + tensors.push((name.clone(), tensor)); + } else { + tensors.push((name.clone(), tensor)); + } + } + }; + } + Ok(tensors) +} + pub fn round_by_factor(num: u32, factor: u32) -> u32 { let round = (num as f32 / factor as f32).round() as u32; round * factor @@ -490,3 +721,40 @@ pub fn get_default_save_dir() -> Option { path.to_string_lossy().to_string() }) } + +pub async fn download_model( + model_id: &str, + save_dir: &str, + max_retries: u32, +) -> anyhow::Result<()> { + let mut attempts = 0u32; + loop { + attempts += 1; + println!( + "Attempting to download model (attempt {}/{})", + attempts, max_retries + ); + + match ModelScope::download(model_id, save_dir).await { + Ok(()) => { + println!("Model downloaded successfully"); + return Ok(()); + } + Err(e) => { + if attempts >= max_retries { + return Err(anyhow::anyhow!( + "Failed to download model after {} attempts. Last error: {}", + max_retries, + e + )); + } + + println!( + "Download failed (attempt {}): {}. Retrying in 2 seconds...", + attempts, e + ); + sleep(Duration::from_secs(2)).await; + } + } + } +} diff --git a/src/utils/tensor_utils.rs b/src/utils/tensor_utils.rs index 3ff8b42..ce0a9cf 100644 --- a/src/utils/tensor_utils.rs +++ b/src/utils/tensor_utils.rs @@ -2,7 +2,23 @@ use anyhow::{Result, anyhow}; use candle_core::{D, DType, Device, IndexOp, Tensor, shape::Dim}; use candle_nn::ops::sigmoid; -pub fn mask_filled(on_true: &Tensor, mask: &Tensor, on_false: f32) -> Result { +pub enum PaddingSide { + Left, + Right, +} + +pub fn masked_fill_zeros(hidden_states: &Tensor, mask: &Tensor) -> Result { + // hidden_states: (bs, seq_len, hidden_dim) + // mask: (bs, seq_len) + let on_false = hidden_states.zeros_like()?; + let mask = mask + .unsqueeze(D::Minus1)? + .broadcast_as(hidden_states.shape())?; + let hidden_states = mask.where_cond(&hidden_states, &on_false)?; + Ok(hidden_states) +} + +pub fn attn_masked_fill(on_true: &Tensor, mask: &Tensor, on_false: f32) -> Result { let (mask_seq_len, _) = mask.dims2()?; let (_, _, seq_len, _) = on_true.dims4()?; assert!( @@ -476,6 +492,44 @@ pub fn interpolate_linear_1d( Ok(output) } +pub fn interpolate_nearest_1d(t: &Tensor, target_size: usize) -> Result { + // t: [b, channels, features] + if t.rank() != 3 { + return Err(anyhow::anyhow!( + "Input rank must have equal to 3 dimensions" + )); + } + let shape = t.dims(); + let orig_size = shape[shape.len() - 1]; + if orig_size == target_size { + return Ok(t.clone()); + } + + let (bs, channels, _) = t.dims3()?; + let mut output = Tensor::zeros((bs, channels, target_size), t.dtype(), t.device())?; + let coords = compute_1d_coords(orig_size, target_size, None)?; + + for b in 0..bs { + for c in 0..channels { + let input_slice = t.i((b, c))?; + let mut out_i = Vec::new(); + + for &coord in coords.iter().take(target_size) { + // Nearest neighbor: round to nearest integer coordinate + let nearest_idx = coord.floor() as usize; + let clamped_idx = nearest_idx.min(orig_size - 1); + + let value = input_slice.get(clamped_idx)?; + out_i.push(value); + } + let out_i = Tensor::stack(&out_i, 0)?.unsqueeze(0)?.unsqueeze(0)?; + output = output.slice_assign(&[(b..b + 1), (c..c + 1), (0..target_size)], &out_i)?; + } + } + output = output.contiguous()?; + Ok(output) +} + pub fn interpolate_bilinear( input: &Tensor, target_size: (usize, usize), @@ -909,3 +963,83 @@ pub fn pad_replicate_last_dim(t: &Tensor, pad: (usize, usize)) -> Result } Ok(pad_tensor) } + +pub fn log10(t: &Tensor) -> Result { + Ok(t.log()?.affine(1.0 / 10.0_f64.ln(), 0.0)?) +} + +pub fn z_score_normalize(t: &Tensor, dim: usize) -> Result { + let rank = t.rank(); + if dim >= rank { + return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank))); + } + Ok(t.broadcast_sub(&t.mean_keepdim(dim)?)? + .broadcast_div(&t.var_keepdim(dim)?.sqrt()?)?) +} + +pub fn l2_normalize(t: &Tensor, dim: usize) -> Result { + let rank = t.rank(); + if dim >= rank { + return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank))); + } + let l2_norm = t.sqr()?.sum_keepdim(dim)?.sqrt()?; + Ok(t.broadcast_div(&l2_norm)?) +} + +pub fn l1_normalize(t: &Tensor, dim: usize) -> Result { + let rank = t.rank(); + if dim >= rank { + return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank))); + } + let l1_norm = t.abs()?.sum_keepdim(dim)?; + Ok(t.broadcast_div(&l1_norm)?) +} + +pub fn pool1d(xs: &Tensor, pool_size: usize, ceil_mode: bool, stype: &str) -> Result { + // xs: (bs, c, dim) + // ceil_mode: 是否保留不完整窗口，为true时通过pad实现 + if pool_size == 0 { + return Err(anyhow!("pool_size must be greater than 0")); + } + let (bs, c, dim) = xs.dims3()?; + let xs_reshape = if ceil_mode { + let remain = dim % pool_size; + if remain > 0 { + let pad = pool_size - remain; + let xs_pad = pad_replicate_last_dim(xs, (0, pad))?; + xs_pad.reshape((bs, c, (), pool_size))? + } else { + xs.reshape((bs, c, (), pool_size))? + } + } else { + let remain = dim % pool_size; + if remain > 0 { + let xs_del = xs.narrow(D::Minus1, 0, dim - remain)?; + xs_del.reshape((bs, c, (), pool_size))? + } else { + xs.reshape((bs, c, (), pool_size))? + } + }; + let xs_pool = match stype { + "avg" => xs_reshape.mean(D::Minus1)?, + "max" => xs_reshape.max(D::Minus1)?, + "min" => xs_reshape.min(D::Minus1)?, + _ => { + return Err(anyhow!( + "unsupported pool type: {}, supported types are: avg, max, min", + stype + )); + } + }; + Ok(xs_pool) +} + +pub fn statistics_pooling(xs: &Tensor, dim: D, keepdim: bool) -> Result { + let mean = xs.mean(dim)?; + let std = xs.var(dim)?.sqrt()?; + let mut stats = Tensor::cat(&[mean, std], D::Minus1)?; + if keepdim { + stats = stats.unsqueeze(dim)?; + } + Ok(stats) +} diff --git a/tests/messy_test.rs b/tests/messy_test.rs index 9a9e150..b653509 100644 --- a/tests/messy_test.rs +++ b/tests/messy_test.rs @@ -1,22 +1,31 @@ // use std::io::Cursor; -use aha::utils::audio_utils::create_hann_window; +use std::time::Instant; + +use aha::utils::{audio_utils::create_hann_window, tensor_utils::interpolate_nearest_1d}; use anyhow::Result; -use candle_core::DType; +use candle_core::{DType, Tensor}; // use symphonia::core::io::MediaSourceStream; #[test] fn messy_test() -> Result<()> { // RUST_BACKTRACE=1 cargo test -F cuda,ffmpeg messy_test -r -- --nocapture let device = &candle_core::Device::Cpu; + let t = Tensor::arange(0.0f32, 40.0, device)?.broadcast_as((1, 40, 40))?; + println!("t: {}", t); + let i_start = Instant::now(); + let t_inter = interpolate_nearest_1d(&t, 20)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in interpolate_nearest_1d is: {:?}", i_duration); + println!("t_inter: {}", t_inter); // let url = "https://sis-sample-audio.obs.cn-north-1.myhuaweicloud.com/16k16bit.mp3"; // let client = reqwest::blocking::Client::new(); // let response = client.get(url).send()?; // let vec_u8 = response.bytes()?.to_vec(); // let mut content = Cursor::new(vec_u8); // let mss = MediaSourceStream::new(Box::new(content), Default::default()); - let window = create_hann_window(400, DType::F32, device)?; - println!("window: {}", window); + // let window = create_hann_window(400, DType::F32, device)?; + // println!("window: {}", window); // let audio_path = "file:///home/jhq/Videos/voice_01.wav"; // let audio_path = "/home/jhq/Videos/zh.mp3"; // let audio_path = "/home/jhq/Videos/zh.mp3"; diff --git a/tests/test_index_tts2.rs b/tests/test_index_tts2.rs new file mode 100644 index 0000000..1fd1745 --- /dev/null +++ b/tests/test_index_tts2.rs @@ -0,0 +1,48 @@ +use std::time::Instant; +use anyhow::Result; + +use aha::models::index_tts2::{generate::IndexTTS2Generate, utils::download_index_tts2_need_model}; +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; + +#[tokio::test] +async fn index_tts2_generate() -> Result<()> { + // RUST_BACKTRACE=1 cargo test -F cuda index_tts2_generate -r -- --nocapture + let save_dir = + aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let _ = download_index_tts2_need_model(Some(&save_dir)).await?; + let model_path = format!("{}/IndexTeam/IndexTTS-2", save_dir); + let message = r#" + { + "model": "index-tts2", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "audio", + "audio_url": + { + "url": "file:///home/jhq/Videos/voice_01.wav" + } + }, + { + "type": "text", + "text": "你好啊" + } + ] + } + ] + } + "#; + let mes: ChatCompletionParameters = serde_json::from_str(message)?; + let i_start = Instant::now(); + let mut voxcpm_generate = IndexTTS2Generate::init(&model_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + + let i_start = Instant::now(); + let generate = voxcpm_generate.generate(mes)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in generate is: {:?}", i_duration); + Ok(()) +} \ No newline at end of file diff --git a/tests/weight_test.rs b/tests/weight_test.rs index a70aa56..f433397 100644 --- a/tests/weight_test.rs +++ b/tests/weight_test.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; -use aha::utils::{find_type_files, get_device}; +use aha::utils::{find_type_files, get_device, read_pth_tensor_info_cycle}; use anyhow::Result; use candle_core::{Device, pickle::read_all_with_key, safetensors}; use candle_nn::VarBuilder; @@ -197,3 +197,41 @@ fn qwen3_weight() -> Result<()> { println!("model_list: {:?}", model_list); Ok(()) } + +#[test] +fn index_tts2_weight() -> Result<()> { + let save_dir: String = + aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/IndexTeam/IndexTTS-2/", save_dir); + let s2mel_path = model_path+ "/s2mel.pth"; + // let wac2vec2_path = model_path+ "/wav2vec2bert_stats.pt"; + // let model_path = format!("{}/iic/speech_campplus_sv_zh-cn_16k-common/", save_dir); + // let campplus_path = model_path+ "/campplus_cn_common.bin"; + // let model_list = find_type_files(&model_path, "safetensors")?; + let model_list = vec![s2mel_path]; + // let mut dict_to_hashmap = HashMap::new(); + // let mut dtype = candle_core::DType::F32; + for m in model_list { + // let dict = read_all_with_key(m, Some("state_dict"))?; + // let dict = read_all_with_key(m, Some("net"))?; + let dict = read_pth_tensor_info_cycle(m, Some("net.cfm"))?; + // dtype = dict[0].1.dtype(); + for (k, v) in dict { + // if k.contains("model") { + // println!("key: {}, tensor shape: {:?}", k, v); + // } + // dict_to_hashmap.insert(k, v); + println!("key: {}, tensor shape: {:?}", k, v); + } + } + // let device = Device::Cpu; + // let semantic_codec_path = save_dir.to_string() + "/amphion/MaskGCT/semantic_codec/model.safetensors" ; + // let model_list = vec![semantic_codec_path]; + // for m in model_list { + // let weights = safetensors::load(m, &device)?; + // for (key, tensor) in weights.iter() { + // println!("=== {} === {:?}", key, tensor.shape()); + // } + // } + Ok(()) +} \ No newline at end of file