From e50c6cb795eb4eb61927660a478a814adb0f72b8 Mon Sep 17 00:00:00 2001 From: Xiaoyang Chen Date: Wed, 14 Oct 2026 18:33:44 -0700 Subject: [PATCH] Split libwaifu into two C libraries and a CLI that only calls them waifu is now libwaifu, a cdylib whose C API runs a model given the path of its manifest. The new waifu-modelmanager crate is libwaifu_modelmanager: the catalogue, fetching, config.toml and what a model's boxes start at. The two share no code and no types, and each header stands alone. - waifu-ffi and waifu-modelmanager's C API crate are folded into their crates (cdylib + rlib). hub.rs and config.rs moved to the model manager; yaml.rs is copied, since both read manifests. - libwaifu: load and describe take a manifest path (ABI 2); describe says only what a request may hold; a draw runs the request as given. New waifu_devices_json. - libwaifu_modelmanager: new manifest_path, directory_json and defaults_json, the last being the kind defaults and the manifest's suggested: block, which moved here from the engine. - waifu-cli: the `waifu` binary, statically linked against both crates and calling nothing but their C APIs. - waifu-tools: the old examples, as binaries. - macOS and Windows apps move to waifu-app/swiftui and waifu-app/wpf and adapt to the API; they behave as before. - IndexTTS's own docs merged into docs/indextts.md; the text normalizer moves to waifu::text_normalizer, shared with CosyVoice3. Co-Authored-By: Claude Opus 5.5 --- .claude/skills/add-anima-model/SKILL.md | 14 +- .claude/skills/add-sdxl-model/SKILL.md | 22 +- .github/workflows/ci.yml | 39 + .github/workflows/release.yml | 8 +- CLAUDE.md | 15 +- CMakeLists.txt | 108 ++- Cargo.lock | 35 +- Cargo.toml | 4 +- README.md | 23 +- docs/TODO.md | 16 +- docs/audio.md | 14 +- ...indextts_features.md => audio_features.md} | 6 +- docs/bigvgan.md | 128 ---- docs/cosyvoice3.md | 10 +- docs/ffi.md | 195 +++-- docs/indextts.md | 306 ++++++-- docs/indextts_emotion.md | 158 ---- docs/indextts_gpt.md | 158 ---- docs/s2mel.md | 148 ---- docs/semantic_codec.md | 100 --- docs/speech.md | 7 +- ...ndextts_normalize.md => text_normalize.md} | 7 +- docs/w2v_bert.md | 67 -- macos/Sources/CWaifu/include/waifu.h | 1 - tools/indextts_emotion_reference.py | 2 +- tools/indextts_gpt_exporter.py | 2 +- tools/indextts_normalize_reference.py | 2 +- tools/semantic_codec_exporter.py | 2 +- tools/w2v_bert_exporter.py | 2 +- {macos => waifu-app/swiftui}/.gitignore | 0 {macos => waifu-app/swiftui}/Package.swift | 16 +- {macos => waifu-app/swiftui}/README.md | 30 +- .../Resources/AppIcon.icon/Assets/picture.png | Bin .../swiftui}/Resources/AppIcon.icon/icon.json | 0 .../swiftui/Sources/CWaifu/include/waifu.h | 1 + .../CWaifu/include/waifu_modelmanager.h | 1 + .../swiftui}/Sources/CWaifu/shim.c | 6 +- .../swiftui}/Sources/Waifu/Launcher.swift | 0 .../swiftui}/Sources/Waifu/MainView.swift | 0 .../swiftui}/Sources/Waifu/Media.swift | 0 .../swiftui}/Sources/Waifu/ModelCard.swift | 0 .../Sources/Waifu/ModelManagerView.swift | 0 .../swiftui}/Sources/Waifu/Models.swift | 2 +- .../swiftui}/Sources/Waifu/OutputPane.swift | 0 .../swiftui}/Sources/Waifu/Session.swift | 0 .../Sources/Waifu/SettingsColumn.swift | 0 .../swiftui}/Sources/Waifu/WaifuApp.swift | 0 .../swiftui}/Sources/WaifuKit/WaifuKit.swift | 216 +++++- .../Tests/WaifuKitTests/WaifuKitTests.swift | 12 +- {macos => waifu-app/swiftui}/build-app.sh | 25 +- {macos => waifu-app/swiftui}/dmg_layout.py | 0 {macos => waifu-app/swiftui}/package-dmg.sh | 10 +- {windows => waifu-app/wpf}/.gitignore | 0 {windows => waifu-app/wpf}/Libwaifu.Gui.sln | 0 .../wpf}/Libwaifu.Gui/App.config | 0 .../wpf}/Libwaifu.Gui/App.xaml | 0 .../wpf}/Libwaifu.Gui/App.xaml.cs | 0 .../wpf}/Libwaifu.Gui/Controls/DropBox.cs | 0 .../wpf}/Libwaifu.Gui/Controls/FormGroup.cs | 0 .../wpf}/Libwaifu.Gui/Controls/SettingRow.cs | 0 .../Libwaifu.Gui/Converters/Converters.cs | 0 .../wpf}/Libwaifu.Gui/Engine/Catalog.cs | 2 +- .../wpf}/Libwaifu.Gui/Engine/Core.cs | 0 .../wpf}/Libwaifu.Gui/Engine/Json.cs | 0 .../wpf}/Libwaifu.Gui/Engine/Library.cs | 65 +- .../wpf}/Libwaifu.Gui/Engine/Machine.cs | 0 .../wpf}/Libwaifu.Gui/Engine/MockEngine.cs | 0 .../wpf}/Libwaifu.Gui/Engine/Models.cs | 0 .../wpf}/Libwaifu.Gui/Engine/Native.cs | 140 +++- .../wpf}/Libwaifu.Gui/Engine/NativeEngine.cs | 74 +- .../wpf}/Libwaifu.Gui/Engine/NativeModels.cs | 121 ++- .../wpf}/Libwaifu.Gui/Engine/Output.cs | 0 .../wpf}/Libwaifu.Gui/Engine/Sound.cs | 0 .../wpf}/Libwaifu.Gui/Engine/Store.cs | 0 .../wpf}/Libwaifu.Gui/Libwaifu.Gui.csproj | 0 .../Libwaifu.Gui/Properties/AssemblyInfo.cs | 0 .../wpf}/Libwaifu.Gui/Theme/Controls.xaml | 0 .../wpf}/Libwaifu.Gui/Theme/Icons.xaml | 0 .../wpf}/Libwaifu.Gui/Theme/Tokens.xaml | 0 .../wpf}/Libwaifu.Gui/ViewModels/Common.cs | 0 .../ViewModels/ImageTaskViewModel.cs | 0 .../ViewModels/ModelManagerViewModel.cs | 0 .../wpf}/Libwaifu.Gui/ViewModels/Shell.cs | 0 .../ViewModels/SpeechTaskViewModel.cs | 0 .../Libwaifu.Gui/ViewModels/TaskViewModel.cs | 0 .../Views/ConversionInspector.xaml | 0 .../Views/ConversionInspector.xaml.cs | 0 .../Views/ConversionTaskView.xaml | 0 .../Views/ConversionTaskView.xaml.cs | 0 .../Libwaifu.Gui/Views/ImageInspector.xaml | 0 .../Libwaifu.Gui/Views/ImageInspector.xaml.cs | 0 .../Libwaifu.Gui/Views/ImageTaskView.xaml | 0 .../Libwaifu.Gui/Views/ImageTaskView.xaml.cs | 0 .../Libwaifu.Gui/Views/InspectorSelector.cs | 0 .../wpf}/Libwaifu.Gui/Views/LogWindow.xaml | 0 .../wpf}/Libwaifu.Gui/Views/LogWindow.xaml.cs | 0 .../wpf}/Libwaifu.Gui/Views/MainWindow.xaml | 0 .../Libwaifu.Gui/Views/MainWindow.xaml.cs | 0 .../Libwaifu.Gui/Views/ModelManagerView.xaml | 0 .../Views/ModelManagerView.xaml.cs | 0 .../wpf}/Libwaifu.Gui/Views/PlayerView.xaml | 0 .../Libwaifu.Gui/Views/PlayerView.xaml.cs | 0 .../Libwaifu.Gui/Views/SpeechInspector.xaml | 0 .../Views/SpeechInspector.xaml.cs | 0 .../Libwaifu.Gui/Views/SpeechTaskView.xaml | 0 .../Libwaifu.Gui/Views/SpeechTaskView.xaml.cs | 0 .../wpf}/Libwaifu.Gui/app.ico | Bin .../wpf}/Libwaifu.Gui/app.manifest | 0 {windows => waifu-app/wpf}/README.md | 22 +- waifu-cli/Cargo.toml | 38 + waifu-cli/build.rs | 90 +++ {waifu/src/cli => waifu-cli/src}/args.rs | 48 +- waifu-cli/src/c.rs | 75 ++ {waifu/src/cli => waifu-cli/src}/draw.rs | 189 +++-- waifu-cli/src/engine.rs | 396 ++++++++++ waifu-cli/src/log.rs | 110 +++ waifu/src/cli/mod.rs => waifu-cli/src/main.rs | 21 +- waifu-cli/src/manager.rs | 310 ++++++++ waifu-cli/src/runtime.rs | 377 +++++++++ {waifu/src/cli => waifu-cli/src}/task.rs | 6 +- {waifu/src/cli => waifu-cli/src}/tui/ask.rs | 6 +- .../src/cli => waifu-cli/src}/tui/devices.rs | 10 +- {waifu/src/cli => waifu-cli/src}/tui/draw.rs | 265 ++++--- {waifu/src/cli => waifu-cli/src}/tui/field.rs | 0 {waifu/src/cli => waifu-cli/src}/tui/files.rs | 2 +- {waifu/src/cli => waifu-cli/src}/tui/mod.rs | 36 +- .../src/cli => waifu-cli/src}/tui/models.rs | 139 ++-- {waifu/src/cli => waifu-cli/src}/tui/tasks.rs | 8 +- waifu-ffi/Cargo.toml | 19 - waifu-modelmanager/Cargo.toml | 59 ++ waifu-modelmanager/build.rs | 62 ++ waifu-modelmanager/cbindgen.toml | 33 + .../include/waifu_modelmanager.h | 180 +++++ waifu-modelmanager/src/c_api.rs | 417 ++++++++++ waifu-modelmanager/src/c_api/support.rs | 289 +++++++ waifu-modelmanager/src/c_api/tests.rs | 187 +++++ {waifu => waifu-modelmanager}/src/config.rs | 0 waifu-modelmanager/src/defaults.rs | 252 ++++++ waifu-modelmanager/src/fetch.rs | 256 +++++++ {waifu => waifu-modelmanager}/src/hub.rs | 88 ++- .../waifu.rs => waifu-modelmanager/src/lib.rs | 37 +- waifu-modelmanager/src/log.rs | 92 +++ .../src/suggested.rs | 60 +- waifu-modelmanager/src/yaml.rs | 429 +++++++++++ waifu-tools/Cargo.toml | 15 + {waifu-ffi => waifu-tools}/build.rs | 37 +- .../src/bin}/cosyvoice3_bench.rs | 2 +- .../src/bin}/fetch_manifests.rs | 7 +- .../src/bin}/generate.rs | 27 +- .../src/bin}/indextts_bench.rs | 2 +- .../src/bin}/inspect.rs | 4 +- .../examples => waifu-tools/src/bin}/load.rs | 2 +- .../src/bin}/memprofile.rs | 10 +- .../src/bin}/parameters.rs | 2 +- .../examples => waifu-tools/src/bin}/speak.rs | 6 +- waifu-tools/src/lib.rs | 41 + waifu/Cargo.toml | 89 +-- waifu/build.rs | 70 +- {waifu-ffi => waifu}/cbindgen.toml | 11 +- {waifu-ffi => waifu}/include/waifu.h | 142 ++-- waifu-ffi/src/lib.rs => waifu/src/c_api.rs | 725 ++++-------------- waifu/src/c_api/support.rs | 310 ++++++++ {waifu-ffi/src => waifu/src/c_api}/tests.rs | 31 +- waifu/src/cosyvoice3/frontend.rs | 4 +- waifu/src/describe.rs | 330 +++----- waifu/src/engine.rs | 341 +++----- waifu/src/indextts/mod.rs | 4 +- waifu/src/lib.rs | 26 +- waifu/src/manifest.rs | 29 +- waifu/src/progress.rs | 61 +- waifu/src/runtime.rs | 159 +--- waifu/src/text_normalizer/chinese.rs | 153 ++++ waifu/src/text_normalizer/english.rs | 150 ++++ .../normalize.rs => text_normalizer/mod.rs} | 455 +---------- waifu/src/text_normalizer/spanish.rs | 184 +++++ waifu/src/yaml.rs | 7 - waifu/tests/indextts_emotion.rs | 4 +- waifu/tests/manifest.rs | 86 +-- waifu/tests/normalize.rs | 2 +- 179 files changed, 6807 insertions(+), 3649 deletions(-) rename docs/{indextts_features.md => audio_features.md} (96%) delete mode 100644 docs/bigvgan.md delete mode 100644 docs/indextts_emotion.md delete mode 100644 docs/indextts_gpt.md delete mode 100644 docs/s2mel.md delete mode 100644 docs/semantic_codec.md rename docs/{indextts_normalize.md => text_normalize.md} (96%) delete mode 100644 docs/w2v_bert.md delete mode 120000 macos/Sources/CWaifu/include/waifu.h rename {macos => waifu-app/swiftui}/.gitignore (100%) rename {macos => waifu-app/swiftui}/Package.swift (61%) rename {macos => waifu-app/swiftui}/README.md (69%) rename {macos => waifu-app/swiftui}/Resources/AppIcon.icon/Assets/picture.png (100%) rename {macos => waifu-app/swiftui}/Resources/AppIcon.icon/icon.json (100%) create mode 120000 waifu-app/swiftui/Sources/CWaifu/include/waifu.h create mode 120000 waifu-app/swiftui/Sources/CWaifu/include/waifu_modelmanager.h rename {macos => waifu-app/swiftui}/Sources/CWaifu/shim.c (79%) rename {macos => waifu-app/swiftui}/Sources/Waifu/Launcher.swift (100%) rename {macos => waifu-app/swiftui}/Sources/Waifu/MainView.swift (100%) rename {macos => waifu-app/swiftui}/Sources/Waifu/Media.swift (100%) rename {macos => waifu-app/swiftui}/Sources/Waifu/ModelCard.swift (100%) rename {macos => waifu-app/swiftui}/Sources/Waifu/ModelManagerView.swift (100%) rename {macos => waifu-app/swiftui}/Sources/Waifu/Models.swift (99%) rename {macos => waifu-app/swiftui}/Sources/Waifu/OutputPane.swift (100%) rename {macos => waifu-app/swiftui}/Sources/Waifu/Session.swift (100%) rename {macos => waifu-app/swiftui}/Sources/Waifu/SettingsColumn.swift (100%) rename {macos => waifu-app/swiftui}/Sources/Waifu/WaifuApp.swift (100%) rename {macos => waifu-app/swiftui}/Sources/WaifuKit/WaifuKit.swift (70%) rename {macos => waifu-app/swiftui}/Tests/WaifuKitTests/WaifuKitTests.swift (88%) rename {macos => waifu-app/swiftui}/build-app.sh (76%) rename {macos => waifu-app/swiftui}/dmg_layout.py (100%) rename {macos => waifu-app/swiftui}/package-dmg.sh (81%) rename {windows => waifu-app/wpf}/.gitignore (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui.sln (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/App.config (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/App.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/App.xaml.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Controls/DropBox.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Controls/FormGroup.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Controls/SettingRow.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Converters/Converters.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/Catalog.cs (99%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/Core.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/Json.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/Library.cs (76%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/Machine.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/MockEngine.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/Models.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/Native.cs (70%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/NativeEngine.cs (90%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/NativeModels.cs (73%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/Output.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/Sound.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Engine/Store.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Libwaifu.Gui.csproj (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Properties/AssemblyInfo.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Theme/Controls.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Theme/Icons.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Theme/Tokens.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/ViewModels/Common.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/ViewModels/ImageTaskViewModel.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/ViewModels/ModelManagerViewModel.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/ViewModels/Shell.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/ViewModels/SpeechTaskViewModel.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/ViewModels/TaskViewModel.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/ConversionInspector.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/ConversionInspector.xaml.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/ConversionTaskView.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/ConversionTaskView.xaml.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/ImageInspector.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/ImageInspector.xaml.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/ImageTaskView.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/ImageTaskView.xaml.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/InspectorSelector.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/LogWindow.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/LogWindow.xaml.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/MainWindow.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/MainWindow.xaml.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/ModelManagerView.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/ModelManagerView.xaml.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/PlayerView.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/PlayerView.xaml.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/SpeechInspector.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/SpeechInspector.xaml.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/SpeechTaskView.xaml (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/Views/SpeechTaskView.xaml.cs (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/app.ico (100%) rename {windows => waifu-app/wpf}/Libwaifu.Gui/app.manifest (100%) rename {windows => waifu-app/wpf}/README.md (88%) create mode 100644 waifu-cli/Cargo.toml create mode 100644 waifu-cli/build.rs rename {waifu/src/cli => waifu-cli/src}/args.rs (91%) create mode 100644 waifu-cli/src/c.rs rename {waifu/src/cli => waifu-cli/src}/draw.rs (60%) create mode 100644 waifu-cli/src/engine.rs create mode 100644 waifu-cli/src/log.rs rename waifu/src/cli/mod.rs => waifu-cli/src/main.rs (92%) create mode 100644 waifu-cli/src/manager.rs create mode 100644 waifu-cli/src/runtime.rs rename {waifu/src/cli => waifu-cli/src}/task.rs (94%) rename {waifu/src/cli => waifu-cli/src}/tui/ask.rs (99%) rename {waifu/src/cli => waifu-cli/src}/tui/devices.rs (98%) rename {waifu/src/cli => waifu-cli/src}/tui/draw.rs (95%) rename {waifu/src/cli => waifu-cli/src}/tui/field.rs (100%) rename {waifu/src/cli => waifu-cli/src}/tui/files.rs (99%) rename {waifu/src/cli => waifu-cli/src}/tui/mod.rs (91%) rename {waifu/src/cli => waifu-cli/src}/tui/models.rs (92%) rename {waifu/src/cli => waifu-cli/src}/tui/tasks.rs (98%) delete mode 100644 waifu-ffi/Cargo.toml create mode 100644 waifu-modelmanager/Cargo.toml create mode 100644 waifu-modelmanager/build.rs create mode 100644 waifu-modelmanager/cbindgen.toml create mode 100644 waifu-modelmanager/include/waifu_modelmanager.h create mode 100644 waifu-modelmanager/src/c_api.rs create mode 100644 waifu-modelmanager/src/c_api/support.rs create mode 100644 waifu-modelmanager/src/c_api/tests.rs rename {waifu => waifu-modelmanager}/src/config.rs (100%) create mode 100644 waifu-modelmanager/src/defaults.rs create mode 100644 waifu-modelmanager/src/fetch.rs rename {waifu => waifu-modelmanager}/src/hub.rs (96%) rename waifu/src/bin/waifu.rs => waifu-modelmanager/src/lib.rs (51%) create mode 100644 waifu-modelmanager/src/log.rs rename {waifu => waifu-modelmanager}/src/suggested.rs (87%) create mode 100644 waifu-modelmanager/src/yaml.rs create mode 100644 waifu-tools/Cargo.toml rename {waifu-ffi => waifu-tools}/build.rs (59%) rename {waifu/examples => waifu-tools/src/bin}/cosyvoice3_bench.rs (97%) rename {waifu/examples => waifu-tools/src/bin}/fetch_manifests.rs (89%) rename {waifu/examples => waifu-tools/src/bin}/generate.rs (83%) rename {waifu/examples => waifu-tools/src/bin}/indextts_bench.rs (98%) rename {waifu/examples => waifu-tools/src/bin}/inspect.rs (94%) rename {waifu/examples => waifu-tools/src/bin}/load.rs (97%) rename {waifu/examples => waifu-tools/src/bin}/memprofile.rs (92%) rename {waifu/examples => waifu-tools/src/bin}/parameters.rs (98%) rename {waifu/examples => waifu-tools/src/bin}/speak.rs (93%) create mode 100644 waifu-tools/src/lib.rs rename {waifu-ffi => waifu}/cbindgen.toml (68%) rename {waifu-ffi => waifu}/include/waifu.h (75%) rename waifu-ffi/src/lib.rs => waifu/src/c_api.rs (55%) create mode 100644 waifu/src/c_api/support.rs rename {waifu-ffi/src => waifu/src/c_api}/tests.rs (90%) create mode 100644 waifu/src/text_normalizer/chinese.rs create mode 100644 waifu/src/text_normalizer/english.rs rename waifu/src/{indextts/normalize.rs => text_normalizer/mod.rs} (54%) create mode 100644 waifu/src/text_normalizer/spanish.rs diff --git a/.claude/skills/add-anima-model/SKILL.md b/.claude/skills/add-anima-model/SKILL.md index 7f97f2a..c81f828 100644 --- a/.claude/skills/add-anima-model/SKILL.md +++ b/.claude/skills/add-anima-model/SKILL.md @@ -251,8 +251,8 @@ local test fixture is the exception and is deliberately not this -- see "Test da ## 4. Verify before uploading, not after ```bash -cargo run --release --example load -- models/.yaml cuda -cargo run --release --example generate -- models/.yaml "" cuda +cargo run --release -p waifu-tools --bin load -- models/.yaml cuda +cargo run --release -p waifu-tools --bin generate -- models/.yaml "" cuda ``` `load` builds the whole model from the manifest and encodes a prompt with it -- weights, @@ -327,7 +327,7 @@ model and not on the pictures. Add `not-for-all-audiences` if the samples warran ## 6. Wire it into the CLI -`waifu/src/cli/hub.rs`, once the model is actually on both hubs and not before -- an entry naming +`waifu-modelmanager/src/hub.rs`, once the model is actually on both hubs and not before -- an entry naming a manifest that is not published is a name that 404s for everyone who types it. - a `Published` entry in `CATALOG`: `name`, `full_name` (what the picker shows), `repo`, @@ -358,11 +358,13 @@ entry makes that paragraph wrong as well as incomplete. ```bash ./build/unittest -cargo test --release --manifest-path waifu/Cargo.toml --features cli -- --test-threads=1 +LIBWAIFU_LIB_DIR="$PWD/build" cargo test --release --manifest-path waifu/Cargo.toml -- --test-threads=1 +cargo test --release -p waifu-modelmanager +LIBWAIFU_LIB_DIR="$PWD/build" cargo test --release -p waifu-cli ``` -`--features cli` is what compiles `src/cli/`, so without it the `hub.rs` tests you just changed do -not run at all. +The `hub.rs` tests you just changed are `waifu-modelmanager`'s, and the picker and the `-m` usage +text read the catalogue through its C API, which is what `waifu-cli`'s tests check. The model tests are `#[ignore]`d and need a real package. **Do not run them unless the user asks** (`CLAUDE.md`); say which were skipped and what runs them. When asked, and for a change that diff --git a/.claude/skills/add-sdxl-model/SKILL.md b/.claude/skills/add-sdxl-model/SKILL.md index 0ef25c0..6b0004d 100644 --- a/.claude/skills/add-sdxl-model/SKILL.md +++ b/.claude/skills/add-sdxl-model/SKILL.md @@ -262,11 +262,11 @@ resolves them without listing them in `names()`. ## 4. Verify before uploading, not after -Two examples, in this order: +Two tools, in this order: ```bash -cargo run --release --example load -- models/.yaml cuda -cargo run --release --example generate -- models/.yaml "" cuda +cargo run --release -p waifu-tools --bin load -- models/.yaml cuda +cargo run --release -p waifu-tools --bin generate -- models/.yaml "" cuda ``` `load` assembles the whole model from the manifest and encodes a prompt with it, which is the @@ -356,8 +356,9 @@ the fine tunes differ: base is `openrail++`, WAI is `cdla-permissive-2.0`, NoobA ## 6. Wire it into the CLI -Only two edits in `waifu/src/cli/hub.rs` -- the picker and the `-m` usage text both build -themselves from `hub::names()`, so nothing else needs telling: +Only two edits in `waifu-modelmanager/src/hub.rs` -- the picker and the `-m` usage text both build +themselves from the catalogue, which libwaifu_modelmanager hands out, so nothing else needs +telling: - add a `Published` entry to `CATALOG` (`name`, `repo`, `manifest`) - add the unversioned alias to `ALIASES` @@ -382,10 +383,12 @@ Then `README.md`: the model table, the `waifu -m` example list, and a `Recent up ```bash ./build/unittest -cargo test --manifest-path waifu/Cargo.toml --features cli -- --test-threads=1 +LIBWAIFU_LIB_DIR="$PWD/build" cargo test --manifest-path waifu/Cargo.toml -- --test-threads=1 +cargo test -p waifu-modelmanager +LIBWAIFU_LIB_DIR="$PWD/build" cargo test -p waifu-cli ``` -Both flags matter -- see traps. To run the SDXL numerical tests you also need `models/` +`--test-threads=1` matters -- see traps. The `hub.rs` tests are `waifu-modelmanager`'s. To run the SDXL numerical tests you also need `models/` populated; see "Test data" below. ## Traps @@ -395,7 +398,7 @@ populated; see "Test data" below. | configure stops on `invalid CUTLASS_ROOT` | `third_party/cutlass` is not vendored, and a CUDA build requires it | `third_party/install_cutlass.sh` | | a CUDA build takes tens of minutes longer than expected | `install_flash_attn.sh` was run out of habit: it compiles every kernel for `sm_80`, `sm_86`, `sm_89`, `sm_90` and `sm_120` | skip it. Since `dd804f1` a build without it breaks nothing -- `stores_and_reads_a_paged_kv_cache` asks `F::paged_attention_available()` and skips itself, and `cuda/flash_attn_test.cc` is only compiled when the option is on. Wanted anyway? Pin `FLASH_ATTN_CUDA_ARCH` to this card's arch | | every `sdxl.rs` test fails with `Aborted: out of memory` | `cargo test` runs one thread per core, each loading a 7 GB model onto one GPU | `-- --test-threads=1` | -| `hub.rs` tests never run | `cli` is a Cargo feature; plain `cargo test` compiles none of `src/cli/` | `--features cli` | +| `hub.rs` tests never run | they are `waifu-modelmanager`'s, and `--manifest-path waifu/Cargo.toml` tests only `waifu` | `cargo test -p waifu-modelmanager` | | CPU verify takes 20+ minutes | no CUDA in the build | build with CUDA | | the export dies on a 106-byte "checkpoint" | a Civitai download without a valid token writes the JSON error body to the output file and `curl` still exits `0` | check it is not JSON (`file .safetensors`) before exporting; verify the sha256 | | a link error in `cargo test` that makes no sense | the disk filled: a debug `cargo test` builds a second dependency tree beside the release one cmake already built | `--release` on the Rust suite, and delete each checkpoint the moment its export finishes | @@ -441,4 +444,5 @@ whole second dependency tree beside the release one CMake already built. On a bo that fills the disk, and a full disk here does not announce itself: the tool cannot even create the file it writes command output to, so *every* command fails before it runs, and the failure before that one is a linker error that looks like a code problem. Run the Rust suite as -`cargo test --release --manifest-path waifu/Cargo.toml --features cli`. +`cargo test --release --manifest-path waifu/Cargo.toml`, and the catalogue's as +`cargo test --release -p waifu-modelmanager`. diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9dd4aee..8bede06 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -27,6 +27,19 @@ jobs: env: LIBWAIFU_LIB_DIR: ${{ github.workspace }}/build + - name: Test (model manager) + run: cargo test -p waifu-modelmanager + + - name: Test (command line) + run: cargo test -p waifu-cli + env: + LIBWAIFU_LIB_DIR: ${{ github.workspace }}/build + + - name: Build (tools) + run: cargo build -p waifu-tools + env: + LIBWAIFU_LIB_DIR: ${{ github.workspace }}/build + macos-metal: name: macOS (CPU + Metal) runs-on: macos-latest @@ -48,6 +61,19 @@ jobs: env: LIBWAIFU_LIB_DIR: ${{ github.workspace }}/build + - name: Test (model manager) + run: cargo test -p waifu-modelmanager + + - name: Test (command line) + run: cargo test -p waifu-cli + env: + LIBWAIFU_LIB_DIR: ${{ github.workspace }}/build + + - name: Build (tools) + run: cargo build -p waifu-tools + env: + LIBWAIFU_LIB_DIR: ${{ github.workspace }}/build + windows-cuda: name: Windows (CUDA) runs-on: windows-2022 @@ -116,3 +142,16 @@ jobs: run: cargo test --manifest-path waifu/Cargo.toml env: LIBWAIFU_LIB_DIR: ${{ github.workspace }}/build + + - name: Test (model manager) + run: cargo test -p waifu-modelmanager + + - name: Test (command line) + run: cargo test -p waifu-cli + env: + LIBWAIFU_LIB_DIR: ${{ github.workspace }}/build + + - name: Build (tools) + run: cargo build -p waifu-tools + env: + LIBWAIFU_LIB_DIR: ${{ github.workspace }}/build diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 6bc2bc5..2999f9b 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -40,7 +40,7 @@ jobs: # The Developer ID certificate goes into a keychain of the job's own, made the default so # that codesign finds the identity and notarytool the profile, and deleted when the job ends. # The notary profile is an App Store Connect API key, under the name a developer's machine - # gives the scripts too (macos/README.md). + # gives the scripts too (waifu-app/swiftui/README.md). - name: Set up signing env: CERT_P12: ${{ secrets.MACOS_CERT_P12 }} @@ -80,7 +80,7 @@ jobs: - name: Build app env: WAIFU_SIGN_IDENTITY: ${{ secrets.MACOS_SIGN_IDENTITY }} - run: macos/build-app.sh + run: waifu-app/swiftui/build-app.sh # Signed, notarized and stapled. Apple's notary service takes minutes as a rule, and can take # an hour or two over an account's first submissions. @@ -89,11 +89,11 @@ jobs: WAIFU_SIGN_IDENTITY: ${{ secrets.MACOS_SIGN_IDENTITY }} WAIFU_NOTARY_PROFILE: waifu-notary run: | - macos/package-dmg.sh + waifu-app/swiftui/package-dmg.sh # Named for what it is, which version, and the one architecture it runs on: the runner # is Apple silicon, and the app is built for nothing else. dmg="libwaifu-gui-$WAIFU_VERSION-macos-arm64.dmg" - cp macos/build/libwaifu-gui.dmg "$dmg" + cp waifu-app/swiftui/build/libwaifu-gui.dmg "$dmg" echo "DMG=$dmg" >> "$GITHUB_ENV" - name: Upload artifact diff --git a/CLAUDE.md b/CLAUDE.md index 9bfa9c3..5b73cce 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -9,14 +9,17 @@ cmake --build build -j$(sysctl -n hw.ncpu) ## Test -Before committing, both fast suites must pass: +Before committing, the fast suites must pass: ```bash ./build/unittest LIBWAIFU_LIB_DIR="$PWD/build" cargo test --manifest-path waifu/Cargo.toml +cargo test -p waifu-modelmanager +LIBWAIFU_LIB_DIR="$PWD/build" cargo test -p waifu-cli +LIBWAIFU_LIB_DIR="$PWD/build" cargo build -p waifu-tools ``` -`LIBWAIFU_LIB_DIR` must be absolute because the Rust build script runs from `waifu/`. +`LIBWAIFU_LIB_DIR` must be absolute because the build scripts run from their crates' directories. ## Licence @@ -55,14 +58,14 @@ cargo test --release --manifest-path waifu/Cargo.toml --no-fail-fast \ cargo test --release --manifest-path waifu/Cargo.toml --no-fail-fast \ --test cosyvoice3 -- --ignored --test-threads=1 -# waifu/src/indextts/, tools/indextts_* (campplus.rs and features.rs: run cosyvoice3's too) +# waifu/src/indextts/, waifu/src/text_normalizer/, tools/indextts_* +# (campplus.rs, features.rs and text_normalizer/: run cosyvoice3's too) # needs models/indextts25.yaml and the two recordings tools/indextts_reference_voices.py writes cargo test --release --manifest-path waifu/Cargo.toml --no-fail-fast \ --test indextts --test indextts_gpt -- --ignored --test-threads=1 ``` - Keep `--no-fail-fast` and `--test-threads=1`; each test loads a full model onto the GPU. -- Do not add `--features cli`. - Run only one GPU job at a time. Check `nvidia-smi` before diagnosing OOM or timing failures. - No individual test may exceed three minutes. Reuse model fixtures across tests instead of loading a model per test. The harness gives every test its own thread, so a thread-local @@ -86,6 +89,10 @@ cargo test --release --manifest-path waifu/Cargo.toml --no-fail-fast \ removed from the runtime and the writer: a package that has them no longer loads and has to be exported again. A ControlNet package must use the same format as the model it steers. +## Apps + +Do not change the behavior of `waifu-app/swiftui` or `waifu-app/wpf` unless the user asks. + ## Git - Never modify, commit to, or push directly to `main`. Use a non-`main` branch and a pull request. diff --git a/CMakeLists.txt b/CMakeLists.txt index 0e89385..5f4aa08 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -628,60 +628,94 @@ file(GENERATE # CMake drives the Rust build too, so a plain `cmake --build build` is the whole build: no # separate `cargo build` step. LIBWAIFU_LIB_DIR is set explicitly so this works regardless of where -# the build directory was configured, matching what waifu/build.rs falls back to on its own. +# the build directory was configured, matching what the build scripts fall back to on their own. # # CARGO_TARGET_DIR points inside the build directory as well, so that a build writes nothing # outside it and two build directories configured differently do not fight over one cache. Cargo -# still puts the binary under a profile directory of its own, so the last step copies it up beside -# the C++ executables, where someone who built in `build` expects to find it. +# still puts what it builds under a profile directory of its own, so each target copies it up +# beside the C++ executables, where someone who built in `build` expects to find it. find_program(CARGO_EXECUTABLE cargo REQUIRED) set(waifu_CARGO_DIR "${CMAKE_BINARY_DIR}/cargo") -set(waifu_CLI "${waifu_CARGO_DIR}/release/waifu${CMAKE_EXECUTABLE_SUFFIX}") -# `--features cli` is what builds the `waifu` binary: it is off by default so that using the crate -# as a library does not drag in a terminal library. -set(waifu_CARGO_FEATURES "cli") -add_custom_target(waifu-cli ALL - COMMAND ${CMAKE_COMMAND} -E env "LIBWAIFU_LIB_DIR=${CMAKE_BINARY_DIR}" - "CARGO_TARGET_DIR=${waifu_CARGO_DIR}" - ${CARGO_EXECUTABLE} build -p waifu --features ${waifu_CARGO_FEATURES} --release - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${waifu_CLI}" "${CMAKE_BINARY_DIR}/waifu${CMAKE_EXECUTABLE_SUFFIX}" - WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} - DEPENDS flint "${CMAKE_BINARY_DIR}/flint_link_flags.txt" ${flint_CARGO_DEPENDS} - COMMENT "cargo build -p waifu --features ${waifu_CARGO_FEATURES} --release" - VERBATIM) +set(waifu_CARGO_OUT "${waifu_CARGO_DIR}/release") -# The C API, as a shared library an app links: libwaifu.dylib, libwaifu.so or waifu.dll, beside the -# executables, and waifu.h under include/. See docs/ffi.md. +# The two libraries, and the command line built on their C APIs. See docs/ffi.md. +# +# libwaifu the engine: libwaifu.dylib, libwaifu.so or waifu.dll +# libwaifu_modelmanager the catalogue and fetching, with no engine and no flint behind it # -# Cargo names the library after its crate, waifu_ffi, so the copy gives it the name its install -# name says it has. Its own build script puts the Metal kernels into it and writes the header. +# Each is copied beside the executables with its header under include/. The two share nothing, +# headers included. On Windows the import library goes too, which is what a program linking the +# DLL is linked against. if (APPLE) - set(waifu_FFI_BUILT "${waifu_CARGO_DIR}/release/libwaifu_ffi.dylib") - set(waifu_FFI "${CMAKE_BINARY_DIR}/libwaifu.dylib") + set(waifu_LIB_PREFIX "lib") + set(waifu_LIB_SUFFIX ".dylib") elseif (WIN32) - set(waifu_FFI_BUILT "${waifu_CARGO_DIR}/release/waifu_ffi.dll") - set(waifu_FFI "${CMAKE_BINARY_DIR}/waifu.dll") + set(waifu_LIB_PREFIX "") + set(waifu_LIB_SUFFIX ".dll") else() - set(waifu_FFI_BUILT "${waifu_CARGO_DIR}/release/libwaifu_ffi.so") - set(waifu_FFI "${CMAKE_BINARY_DIR}/libwaifu.so") + set(waifu_LIB_PREFIX "lib") + set(waifu_LIB_SUFFIX ".so") endif() -add_custom_target(waifu-ffi ALL + +# waifu_copy_commands(
): the commands that copy a library +# cargo built, and its header, up beside the executables. +function(waifu_copy_commands out library header_dir header) + set(commands + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${waifu_CARGO_OUT}/${waifu_LIB_PREFIX}${library}${waifu_LIB_SUFFIX}" + "${CMAKE_BINARY_DIR}/${waifu_LIB_PREFIX}${library}${waifu_LIB_SUFFIX}") + if (WIN32) + list(APPEND commands + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${waifu_CARGO_OUT}/${library}.dll.lib" "${CMAKE_BINARY_DIR}/${library}.dll.lib") + endif() + list(APPEND commands + COMMAND ${CMAKE_COMMAND} -E make_directory "${CMAKE_BINARY_DIR}/include" + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${CMAKE_CURRENT_SOURCE_DIR}/${header_dir}/${header}" "${CMAKE_BINARY_DIR}/include/${header}") + set(${out} ${commands} PARENT_SCOPE) +endfunction() + +# The engine. Its own build script reads flint_link_flags.txt, puts the Metal kernels into it and +# writes the header. +waifu_copy_commands(waifu_ENGINE_COPY waifu waifu/include waifu.h) +add_custom_target(waifu-engine ALL COMMAND ${CMAKE_COMMAND} -E env "LIBWAIFU_LIB_DIR=${CMAKE_BINARY_DIR}" "CARGO_TARGET_DIR=${waifu_CARGO_DIR}" - ${CARGO_EXECUTABLE} build -p waifu-ffi --release - COMMAND ${CMAKE_COMMAND} -E copy_if_different "${waifu_FFI_BUILT}" "${waifu_FFI}" - COMMAND ${CMAKE_COMMAND} -E make_directory "${CMAKE_BINARY_DIR}/include" - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${CMAKE_CURRENT_SOURCE_DIR}/waifu-ffi/include/waifu.h" - "${CMAKE_BINARY_DIR}/include/waifu.h" + ${CARGO_EXECUTABLE} build -p waifu --lib --release + ${waifu_ENGINE_COPY} WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} DEPENDS flint "${CMAKE_BINARY_DIR}/flint_link_flags.txt" ${flint_CARGO_DEPENDS} - COMMENT "cargo build -p waifu-ffi --release" + COMMENT "cargo build -p waifu --lib --release" VERBATIM) + +# The model manager. It links nothing native, so it needs nothing CMake built. +waifu_copy_commands(waifu_MODELMANAGER_COPY waifu_modelmanager waifu-modelmanager/include waifu_modelmanager.h) +add_custom_target(waifu-modelmanager ALL + COMMAND ${CMAKE_COMMAND} -E env "CARGO_TARGET_DIR=${waifu_CARGO_DIR}" + ${CARGO_EXECUTABLE} build -p waifu-modelmanager --lib --release + ${waifu_MODELMANAGER_COPY} + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} + COMMENT "cargo build -p waifu-modelmanager --lib --release" + VERBATIM) + +# The command line, `waifu`, which calls the two libraries above through nothing but their C APIs, +# with both linked into it statically: one file, which needs neither library beside it. +set(waifu_CLI "${waifu_CARGO_OUT}/waifu${CMAKE_EXECUTABLE_SUFFIX}") +add_custom_target(waifu-cli ALL + COMMAND ${CMAKE_COMMAND} -E env "LIBWAIFU_LIB_DIR=${CMAKE_BINARY_DIR}" + "CARGO_TARGET_DIR=${waifu_CARGO_DIR}" + ${CARGO_EXECUTABLE} build -p waifu-cli --release + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${waifu_CLI}" "${CMAKE_BINARY_DIR}/waifu${CMAKE_EXECUTABLE_SUFFIX}" + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} + COMMENT "cargo build -p waifu-cli --release" + VERBATIM) + # One cargo at a time in one target directory: two would wait on each other's lock anyway, and -# this way the second reuses what the first built. -add_dependencies(waifu-ffi waifu-cli) +# this way each reuses what the one before built. +add_dependencies(waifu-modelmanager waifu-engine) +add_dependencies(waifu-cli waifu-engine waifu-modelmanager) set(unittest_LIBADD flint_test diff --git a/Cargo.lock b/Cargo.lock index e7376e6..314d2e9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1377,18 +1377,18 @@ dependencies = [ [[package]] name = "indoc" -version = "2.0.7" +version = "2.0.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +checksum = "a37b2691796cffeb8a8cd305ac66e65841559f147f4e63231d0eafa4db5384d1" dependencies = [ "rustversion", ] [[package]] name = "instability" -version = "0.3.13" +version = "0.3.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2bf84e73fa6f27f299dec58e13223cf70db80da872eb921d4f6138342a0eabc8" +checksum = "4c3b5acc1e2fd9375041a388da33d1eb8aed5f7a8c0dd3543e3ea2805adfbe20" dependencies = [ "darling 0.24.1", "indoc", @@ -3419,23 +3419,42 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" name = "waifu" version = "0.1.0" dependencies = [ - "hf-hub", + "cbindgen", "image", "memmap2", - "ratatui", "safetensors", "serde_json", "tokenizers", +] + +[[package]] +name = "waifu-cli" +version = "0.1.0" +dependencies = [ + "image", + "ratatui", + "serde_json", + "waifu", + "waifu-modelmanager", +] + +[[package]] +name = "waifu-modelmanager" +version = "0.1.0" +dependencies = [ + "cbindgen", + "hf-hub", + "serde_json", "toml 1.1.6+spec-1.1.0", "ureq", ] [[package]] -name = "waifu-ffi" +name = "waifu-tools" version = "0.1.0" dependencies = [ - "cbindgen", "waifu", + "waifu-modelmanager", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 2ba4192..4a45996 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,8 @@ [workspace] members = [ "waifu", - "waifu-ffi", + "waifu-cli", + "waifu-modelmanager", + "waifu-tools", ] resolver = "2" diff --git a/README.md b/README.md index 3eee530..115c892 100644 --- a/README.md +++ b/README.md @@ -24,8 +24,9 @@ hardware allows. No API key, no cloud, no queue, and no one else seeing what you ## Supported voices The voices run in the library and through its C API -- which the macOS app speaks with -- rather -than in the `waifu` terminal, which draws. From a checkout, `examples/speak.rs` says a sentence -with either of them in the voice of a recording. +than in the `waifu` terminal, which draws. From a checkout, +`cargo run --release -p waifu-tools --bin speak` says a sentence with either of them in the voice +of a recording. | name | model | published as | |---|---|---| @@ -202,17 +203,17 @@ the defaults where it has one. `Sdxl::generate_reporting` is the same run with a reporter that hears how far along it is and can stop it between steps, which is what the `draw` command is built on. -After completing the build steps below, run the complete example with: +After completing the build steps below, run the complete program with: ```bash -cargo run --release -p waifu --example generate -- \ +cargo run --release -p waifu-tools --bin generate -- \ sdxl.yaml \ "a photo of an astronaut riding a horse on mars" ``` It takes an optional third argument -- `cpu`, `cuda` or `metal` -- and takes the first accelerator this build can reach when that is left out. See -[waifu/examples/generate.rs](waifu/examples/generate.rs) for the complete source. +[waifu-tools/src/bin/generate.rs](waifu-tools/src/bin/generate.rs) for the complete source. ## Build @@ -314,11 +315,13 @@ cmake --build build --target unittest --parallel Run the Rust tests. These read the link flags CMake already wrote out, so they work without re-running `cmake --build` -- just be sure `build/` reflects the latest C++ if you edited a kernel. -`--features cli` is what compiles `waifu/src/cli`, and without it none of the command line is -tested: +The command line, `waifu-cli`, has both libraries linked into it statically and calls them through +their C APIs only: ```bash -cargo test -p waifu --features cli +LIBWAIFU_LIB_DIR="$PWD/build" cargo test -p waifu +cargo test -p waifu-modelmanager +LIBWAIFU_LIB_DIR="$PWD/build" cargo test -p waifu-cli ``` The ignored Rust CUDA integration tests can be run on a CUDA machine with: @@ -354,9 +357,7 @@ cargo test --release -p waifu --no-fail-fast \ Neither flag is optional. `--test-threads=1` keeps one model on the card at a time, where cargo would otherwise start one per core, and `--no-fail-fast` keeps a failure in one test binary from -stopping the ones after it. Leave `--features cli` off these: nothing under `waifu/src/cli` is -reached from `tests/`, and asking for it drags hyper and rustls through a release compile first. -One GPU means one job -- a second checkout running these at the same time reports `Aborted: out of +stopping the ones after it. One GPU means one job -- a second checkout running these at the same time reports `Aborted: out of memory`, which reads as a code failure and is not one. ## License diff --git a/docs/TODO.md b/docs/TODO.md index c145a43..6a8f2f4 100644 --- a/docs/TODO.md +++ b/docs/TODO.md @@ -584,15 +584,11 @@ The Rust test has the gate this entry asked for: `fl_paged_attention_available() the kernels. It matters more than it did, since `WITH_FLASH_ATTN` is now off by default. -## The Rust suite skips the CLI, and cannot run the SDXL tests in parallel +## The Rust suite cannot run the SDXL tests in parallel -Two things about `cargo test --manifest-path waifu/Cargo.toml`, which is the command CLAUDE.md and -`.github/workflows/ci.yml` both give: - -**It compiles none of `src/cli/`.** `waifu/Cargo.toml:23` puts the `cli` module behind a feature, -so every test in `cli/hub.rs`, `cli/args.rs` and `cli/picker.rs` is silently absent -- 153 of them -against 30 without. A broken model catalog entry passes CI. `--features cli` is the fix, and CI -does not pass it. +About `cargo test --manifest-path waifu/Cargo.toml`, which is the command CLAUDE.md and +`.github/workflows/ci.yml` both give. (The command line and the catalogue are crates of their own +now, `waifu-cli` and `waifu-modelmanager`, and both have a step of their own in CI.) **Its default parallelism cannot work on one GPU.** `cargo test` runs one thread per core, and every test in `waifu/tests/sdxl.rs` calls `model()`, which loads a 7 GB package onto the device. @@ -610,10 +606,10 @@ like `refuses_a_size_it_cannot_work_in`, which touch no reference data, fail wit from a line that only loads the model, and it looks like a broken package rather than a harness that asked for too much at once. -Nothing documents either flag. Together the command that actually runs everything is +Nothing documents the flag. The command that actually runs everything is ```bash -cargo test --manifest-path waifu/Cargo.toml --features cli -- --include-ignored --test-threads=1 +cargo test --manifest-path waifu/Cargo.toml -- --include-ignored --test-threads=1 ``` diff --git a/docs/audio.md b/docs/audio.md index 3c3840e..d7d965d 100644 --- a/docs/audio.md +++ b/docs/audio.md @@ -198,16 +198,16 @@ is in the file.) These are the primitives. What sits on top of them is now written, and what joins them up is not. -All six of IndexTTS-2.5's models -- about 5.5 GB between them -- are here, each with a document of -its own and a test against the implementation it came from: +All six of IndexTTS-2.5's models -- about 5.5 GB between them -- are here, each with a test against +the implementation it came from, and described in [indextts.md](indextts.md): | piece | | | --- | --- | -| GPT backbone, 1280d × 24L | `waifu::indextts::gpt`, and [indextts_gpt.md](indextts_gpt.md) -- with the generation loop on top of it | -| semantic codec, 8192 entries | `waifu::indextts::semantic_codec`, and [semantic_codec.md](semantic_codec.md) | -| S2Mel, flow-matching DiT 13L × 512d | `waifu::indextts::s2mel`, and [s2mel.md](s2mel.md) | -| BigVGAN vocoder | `waifu::indextts::bigvgan`, and [bigvgan.md](bigvgan.md) | -| w2v-bert-2.0 | `waifu::indextts::w2v_bert`, and [w2v_bert.md](w2v_bert.md) | +| GPT backbone, 1280d × 24L | `waifu::indextts::gpt` -- with the generation loop on top of it | +| semantic codec, 8192 entries | `waifu::indextts::semantic_codec` | +| S2Mel, flow-matching DiT 13L × 512d | `waifu::indextts::s2mel` | +| BigVGAN vocoder | `waifu::indextts::bigvgan` | +| w2v-bert-2.0 | `waifu::indextts::w2v_bert` | | CAMPPlus | `waifu::indextts::campplus`, and [campplus.md](campplus.md) | What remains is not a model. It is the emotion conformer and perceiver that produce `emo_vec` diff --git a/docs/indextts_features.md b/docs/audio_features.md similarity index 96% rename from docs/indextts_features.md rename to docs/audio_features.md index 45d2550..8e40867 100644 --- a/docs/indextts_features.md +++ b/docs/audio_features.md @@ -1,6 +1,8 @@ -# IndexTTS-2.5's audio front end +# The audio front end -What a recording becomes before any of the models sees it. +What a recording becomes before any of the models sees it: `waifu::indextts::features`, written for +[IndexTTS-2.5](indextts.md) and used by [CosyVoice3](cosyvoice3.md) too -- for CAMPPlus's energies +and for resampling. ```rust use waifu::indextts::features; diff --git a/docs/bigvgan.md b/docs/bigvgan.md deleted file mode 100644 index 0c92065..0000000 --- a/docs/bigvgan.md +++ /dev/null @@ -1,128 +0,0 @@ -# BigVGAN - -The last model in a speech pipeline and the only one whose output is audio: a mel spectrogram in, -a waveform out. - -```rust -use waifu::indextts::bigvgan::{BigVgan, BigVganConfig}; - -let vocoder = BigVgan::build( - BigVganConfig::v2_22khz_80band_256x(), "", &weights, Device::Cuda, DType::Float, -)?; - -let wave = vocoder.forward(&mel)?; // (1, 80, frames) -> (1, 1, frames * 256) -``` - -`nvidia/bigvgan_v2_22khz_80band_256x` is the release IndexTTS-2.5 fetches on first run: 112 M -parameters, 80 mel bands, 22.05 kHz out, and one mel frame for every 256 samples. Nothing about -it is IndexTTS's -- it is NVIDIA's, published separately under the MIT licence and shared by -several speech models. It sits under `indextts` because that is the only pipeline here that -runs it; the day a second one does, it moves back up beside [`waifu::audio`](audio.md). - -## It adds no operator - -Six stages of transposed convolution, each followed by three residual blocks whose outputs are -averaged, and a hundred and nine anti-aliased snakes. Every one of those is a function in -[`waifu::audio`](audio.md), and every one of those is `conv2d`, `matmul` and arithmetic. So this -model ran on all three backends the day it was written, and a card that can draw can also speak. - -What it did need was two things that module did not have, both of them in the activation: -`Padding::Replicate`, and the Kaiser-windowed sinc pair `upsample1d` and `downsample1d`. - -## The activation is the model - -A BigVGAN is a HiFi-GAN with the activation replaced, and everything interesting is in the -replacement. - -**It is periodic.** `x + sin(alpha x)^2 / beta` -- a *snake*. A ReLU has to learn that speech is -made of things that repeat; this is told. `alpha` sets the frequency and `beta` the depth, one -pair per channel, and they are trained. - -**It is anti-aliased.** A periodic function applied at the signal's own rate generates components -above half that rate, and there is nowhere for them to go but back down into the audible band, -folded. So each activation upsamples by two, applies the snake at the higher rate where those -components fit, and filters on the way back down. That is the AMP -- anti-aliased -multi-periodicity -- in the block's name. - -The resampling is the fiddly part, and it is fiddly in a way that is entirely about edges: - -- The filter is a twelve-tap Kaiser-windowed sinc at a cutoff of a quarter and a transition width - of 0.3, and the Kaiser beta is derived from those by the standard design formula. It is not - [`resample_kernel`](audio.md), which is the Hann-windowed one that converts between two sample - rates. The two do the same job with different windows and are not interchangeable. -- Padding is by *replication*, not zeros and not a mirror -- and by enough that the taper a - transposed convolution leaves on each end lands entirely in the padding, which is then cut off. - The upsampling pads by `kernel / ratio - 1` and trims `pad * ratio + (kernel - ratio) / 2` off - the front and one more off the back when the kernel is odd. Those asymmetric halves are the - reference's, and getting one of them wrong shifts the whole waveform by a sample. -- The reference convolves with `groups = C` and the same taps in every group. That is not a - grouped convolution: reading `(N, C, L)` as `(N * C, 1, L)` makes each channel an item in a - batch and one ungrouped convolution does all of them. Which is why neither this nor the - transposed convolution needed the group count `audio::conv_transpose1d` does not have. - -### What it costs - -This is where the vocoder's time goes, and it is worth being plain about it. - -The snake is written as arithmetic -- about six passes over the tensor plus a copy -- where a -fused kernel would read it once and write it once. There are eighteen of them per upsampling -stage, on tensors that have already been lengthened, and each one runs at *twice* the rate -because of the anti-aliasing. NVIDIA ships a CUDA kernel for exactly this -(`alias_free_activation/cuda`), which is what `use_cuda_kernel=True` reaches, and the fact that -they bothered is the measurement. - -So if this model turns out to be slow, the activation is where to look first, and -`audio::snake` together with the two resamplings is the thing to replace. Replacing them changes -nothing above: what a caller sees is `BigVgan::forward`. - -## The weights are the checkpoint's, folded - -Every convolution in the reference is wrapped in `weight_norm`, which stores a direction and a -magnitude and multiplies them at each call. `IndexTTS2.__init__` calls `remove_weight_norm()` as -it loads the vocoder, so the folded weight is what the reference actually runs; -`tools/bigvgan_exporter.py` writes that product and the graph reads it. - -Names are otherwise untouched -- `conv_pre`, `ups.0.0`, `resblocks.3.convs1.1`, `conv_post`, -`activations.2.act.alpha` -- so a tensor in the file and a load in the graph can be compared -directly when one of them is wrong. - -The snake's `alpha` and `beta` are stored as logarithms and raised where they are read, which is -what the reference does and what keeps the package a copy rather than a transformation. It costs -one node over a vector as long as the channel count. - -## How it is checked - -`waifu/tests/bigvgan.rs` runs against numbers that came out of NVIDIA's implementation, not out -of a second reading of the paper. `tools/bigvgan_reference.py` downloads `bigvgan.py`, builds a -small model -- 8 mel bands, two upsamplings, 24 samples out -- and prints what it produced as -Rust constants. - -Small, but not simplified: it has both kinds of convolution, four AMP blocks, dilations, and -seventeen anti-aliased snakes, which is every distinct thing the 112 M parameter release does. And -it is in the *fast* suite, because neither side reads a weight from the other -- both fill every -parameter from a hash of its own name, in float64 rounded once, so the two models are equal bit -for bit with no file passing between them. A vocoder that stops matching BigVGAN is noticed in -the same minute it stops, rather than the next time someone runs the ignored tests. - -The tests that do need the real 460 MB checkpoint are `#[ignore]`d in the same file, and what -they compare is a chirp's mel spectrogram through the whole thing: - -```bash -.venv/bin/python tools/bigvgan_exporter.py \ - -output models/bigvgan-22khz-80band.safetensors \ - -test_output models/bigvgan-22khz-80band_test.safetensors - -cargo test --release --manifest-path waifu/Cargo.toml --test bigvgan -- --ignored -``` - -## What is not here - -`resblock: "2"`, the cheaper AMP block with one convolution per dilation instead of two. No -released BigVGAN v2 uses it, and a configuration asking for it is refused rather than quietly run -as the other kind. - -A kernel for the snake, and the fused CUDA activation that goes with it. See what it costs, above. - -And the five models above this one. `docs/audio.md` lists what IndexTTS-2.5 still needs; with the -vocoder written, what is left is the GPT backbone, the semantic codec, S2Mel, the Qwen3 emotion -model, and the two encoders that are not in the model's own repository. diff --git a/docs/cosyvoice3.md b/docs/cosyvoice3.md index 2f297b7..03e822d 100644 --- a/docs/cosyvoice3.md +++ b/docs/cosyvoice3.md @@ -6,8 +6,8 @@ under Apache 2.0. ```bash waifu -m cosyvoice # the text2speech page -cargo run --release --example speak -- models/cosyvoice3.yaml voice.wav "你好。" out.wav -cargo run --release --example speak -- models/cosyvoice3.yaml voice.wav "你好。" out.wav 7 "录音里说的话。" +cargo run --release -p waifu-tools --bin speak -- models/cosyvoice3.yaml voice.wav "你好。" out.wav +cargo run --release -p waifu-tools --bin speak -- models/cosyvoice3.yaml voice.wav "你好。" out.wav 7 "录音里说的话。" ``` ```rust @@ -57,9 +57,9 @@ The exporter checks the written `tokenizer.json` against the tokenizer upstream Tokens and mel are then cut to the same length, two frames a token. -**Saying, once per sentence.** The text is normalized and split into pieces of 60 to 80 -characters (Chinese) or tokens (anything else), as upstream's `text_normalize` does -- see -[`frontend.rs`](../waifu/src/cosyvoice3/frontend.rs) -- and each piece goes through: +**Saying, once per sentence.** The text is [normalized](text_normalize.md) and split into pieces +of 60 to 80 characters (Chinese) or tokens (anything else), as upstream's `text_normalize` does -- +see [`frontend.rs`](../waifu/src/cosyvoice3/frontend.rs) -- and each piece goes through: 1. [The language model](../waifu/src/cosyvoice3/lm.rs): `[sos] text [task] prompt-tokens` in, speech tokens out, by repetition-aware nucleus sampling (top-p 0.8, top-k 25, a token said in diff --git a/docs/ffi.md b/docs/ffi.md index 3439196..9eb04ec 100644 --- a/docs/ffi.md +++ b/docs/ffi.md @@ -1,6 +1,9 @@ # A C API for libwaifu -Status: design, not implemented. +Status: implemented. The headers -- `waifu/include/waifu.h` and +`waifu-modelmanager/include/waifu_modelmanager.h` -- are generated from the code and are the last +word; the listing below is the design, kept for its comments, and libwaifu_modelmanager's half of +it is written with the names libwaifu's has rather than its own (see below). The macOS app talks to `waifu webui` over HTTP: it starts the program as a child, waits for it to print its address, and polls it twice a second. This replaces that with a library the app links: @@ -10,19 +13,58 @@ same header serves both, and any other language that can call C. ## Shape ``` -waifu/ the library (rlib), with what the C API needs moved out of `cli` - engine.rs Engine: one thread that holds a model and runs jobs on it -- safe Rust, - tested with cargo test, and what webui's worker can become later -waifu-ffi/ cdylib: extern "C" over Engine and the catalog, nothing else - build.rs the link flags - waifu.h generated by cbindgen -macos/ Swift: a module map over waifu.h, and ~300 lines making it async/await -dotnet/ (later) C#: [LibraryImport] over the same library, ~300 lines making it Task-based +waifu/ libwaifu (cdylib + rlib): the engine, describing a model, the machine + src/c_api.rs extern "C" over the engine, nothing else + include/waifu.h generated by cbindgen +waifu-modelmanager/ libwaifu_modelmanager (cdylib + rlib): the catalogue, fetching, config.toml + src/c_api.rs extern "C" over it, nothing else + include/waifu_modelmanager.h +waifu-cli/ the `waifu` command line: both libraries linked in statically, called + through their C APIs only +waifu-tools/ programs for working on the libraries: checking an export, timing a model +waifu-app/ + swiftui/ Swift: a module map over the headers, and ~300 lines making it async/await + wpf/ C#: [DllImport] over the same libraries, Task-based ``` -Everything with logic in it is Rust behind a safe API. `waifu-ffi` only converts: C strings to +Everything with logic in it is Rust behind a safe API. `c_api.rs` only converts: C strings to `&str`, structs to option structs, results to borrowed C views, panics to status codes. +There are two libraries, and neither links the other. **libwaifu_modelmanager** lists, fetches +and removes published models, and finds their manifests; it links no engine and no flint, so an +app can manage models with it alone. **libwaifu** is the engine: it is handed a model as the path +of its manifest, and knows nothing about which models are published or where they are fetched +from. Loading a published model is two calls, one to each: + +```c +// Fetch it, if the catalogue says it is not all here, then ask where its manifest is. +WaifuModelFetch *fetch = waifu_modelmanager_fetch_async("sdxl:base", ..., on_complete); +// ... on_complete with WAIFU_OK ... +char *manifest = waifu_modelmanager_manifest_path("sdxl:base"); +waifu_engine_load_async(engine, manifest, ...); +waifu_modelmanager_string_free(manifest); +``` + +What else that means for a caller: + +- Each header stands alone: it includes nothing but `stdbool.h`, `stddef.h` and `stdint.h`, and + declares every type it uses. The two share none -- libwaifu_modelmanager's carry its own prefix, + as its functions do: `WaifuModelManagerStatusCode` (`WAIFU_MODELMANAGER_OK`, ...), + `WaifuModelManagerStatus`, `WaifuModelManagerLogLevel`/`LogEvent`, `WaifuModelManagerProgressEvent` + (a download's: no stage, and which `hub`), `WaifuModelManagerCompletionEvent`, and the + `WaifuModelFetch` handle -- so both headers can be included together, and either library's can + change without the other. +- Each library keeps its own last error, string allocator, and log callback and level: + `waifu_last_error`/`waifu_modelmanager_last_error`, + `waifu_string_free`/`waifu_modelmanager_string_free`, and so on. An app sets the log callback on + both. +- Describing a model takes the models' own code to say what they suggest, so it is libwaifu's: + `waifu_describe_json(path)`. A published model's manifest is described before its weights are + fetched -- one fetched on its own, or bundled with the app -- by asking the model manager for + its path first. +- The stand-in voice, `"tones"`, is libwaifu's own and has no manifest: it is loaded and + described by that name. + ## What has to move out of `cli` first The C API needs things that only build with the `cli` feature today. They move into the library, @@ -80,7 +122,7 @@ pixels. - **Results out** are handed to a callback as views that are valid only until the callback returns. The caller copies what it keeps: Swift into `Data` and `[Float]`, .NET into managed arrays. Nothing out needs freeing except strings returned by synchronous calls, which are - freed with `waifu_string_free`. + freed with the `_string_free` of the library that returned them. - **Every callback takes `(void *user_data, const Waifu…Event *event)`**, and nothing else: what it is told is one struct. The event and everything it points to are valid only until the callback returns. Fields added to an event later go at the end; the library writes events and @@ -92,10 +134,11 @@ pixels. - **Asynchronous functions end in `_async`.** They return at once, and report through callbacks: `on_progress` any number of times, then `on_complete`. - **Errors.** Synchronous calls return a `WaifuStatusCode` and leave the message in - `waifu_last_error()`, which is per thread. Asynchronous calls fail in one of two ways: + `waifu_last_error()` (`waifu_modelmanager_last_error()` for the model manager), which is per + thread. Asynchronous calls fail in one of two ways: - **At the call**, for what can be told without running anything: a NULL, a string that is not UTF-8, a `struct_size` too small, a freed engine. The function returns `0` (a NULL handle - for `waifu_modelmanager_fetch_async`) with the message in `waifu_last_error()`, and **no + for `waifu_modelmanager_fetch_async`) with the message in the library's last error, and **no callback is ever called** -- not `on_progress`, not `on_complete`. - **While running**, for everything else: a download that fails, a model that cannot be read or cannot do what was asked, a job for a model of the wrong kind (only known when the job's @@ -124,14 +167,14 @@ pixels. #include #include -#define WAIFU_ABI_VERSION 1 +#define WAIFU_ABI_VERSION 2 // 2: load and describe take a manifest's path, not a published name uint32_t waifu_abi_version(void); typedef enum { WAIFU_OK = 0, WAIFU_CANCELLED = 1, WAIFU_ERR_INVALID_ARGUMENT = 2, // a NULL, a string that is not UTF-8, a size out of range - WAIFU_ERR_UNKNOWN_MODEL = 3, // neither a published name nor a manifest that exists + WAIFU_ERR_UNKNOWN_MODEL = 3, // no published model by that name, or no manifest at that path WAIFU_ERR_FETCH = 4, // the download failed WAIFU_ERR_MODEL = 5, // the model could not be read, or cannot do what was asked WAIFU_ERR_NOT_LOADED = 6, // a job, with no model, or a model of the wrong kind @@ -181,8 +224,8 @@ char *waifu_machine_json(void); // -- progress -------------------------------------------------------------------------------- typedef enum { - WAIFU_STAGE_FETCHING, // file, done/total in bytes, part/parts - WAIFU_STAGE_READING, // reading weights onto the device: no fraction + // 0 was FETCHING: downloads are libwaifu_modelmanager's, with an event of their own + WAIFU_STAGE_READING = 1, // reading weights onto the device: no fraction WAIFU_STAGE_ENCODING, // the prompt, or the picture img2img starts from WAIFU_STAGE_DRAWING, // done/total in steps WAIFU_STAGE_DECODING, @@ -195,8 +238,8 @@ typedef struct { WaifuStage stage; double fraction; // 0..1 of the whole job, or -1 where nothing can say uint64_t done, total; // what stage counts; 0, 0 where it counts nothing - uint32_t part, parts; // FETCHING only: which file of how many - const char *file; // FETCHING only, else NULL + uint32_t part, parts; // reserved: 0, 0 -- they were a download's + const char *file; // reserved: NULL const char *words; // "step 3 of 8", for a status line double seconds; // since the job started } WaifuProgressEvent; @@ -212,25 +255,35 @@ typedef void (*WaifuCompleteCallback)(void *user_data, const WaifuCompletionEven // -- the model manager: what is published, what is here, and fetching it -------------------- // -// None of these needs an engine, or loads anything onto a device. +// libwaifu_modelmanager, in waifu_modelmanager.h. None of these needs an engine, or loads +// anything onto a device. Its own abi_version, last_error, string_free, set_log_callback and +// set_log_level are as libwaifu's below, prefixed waifu_modelmanager_. -// JSON documents; schemas below. Free with waifu_string_free. NULL on failure. +// A JSON document; schema below. Free with waifu_modelmanager_string_free. NULL on failure. char *waifu_modelmanager_catalog_json(void); -char *waifu_modelmanager_describe_json(const char *model); // a published name or a manifest path char *waifu_modelmanager_directory(void); +// {"directory", "said": "environment"|"config"|"default", "default", "config_file"} +char *waifu_modelmanager_directory_json(void); WaifuStatusCode waifu_modelmanager_set_directory(const char *path); // saved to config.toml +// The manifest of a published model here -- the fetched one, else a bundled one -- for libwaifu; +// a manifest's own path back as it is. NULL where there is none here. Only a model the catalogue +// calls `cached` has its weights beside it. +char *waifu_modelmanager_manifest_path(const char *name); +// What a screen should start a model's boxes at -- the kind's defaults with the manifest's +// `suggested:` block over them -- for the path of a manifest, or "tones". Schema below. +char *waifu_modelmanager_defaults_json(const char *model); WaifuStatusCode waifu_modelmanager_remove(const char *name); // deletes a downloaded package // Reads published models' manifests out of `directory` where the download directory has none: -// an app built with every model's manifest (waifu/examples/fetch_manifests.rs fetches them) can -// say what each model suggests before anything is downloaded. Laid out as the download directory -// is, and only read. NULL or "" for none; not saved. +// an app built with every model's manifest (waifu-tools' fetch_manifests fetches them) can say +// what each model suggests before anything is downloaded. Laid out as the +// download directory is, and only read. NULL or "" for none; not saved. WaifuStatusCode waifu_modelmanager_set_bundled_manifests(const char *directory); typedef struct WaifuModelFetch WaifuModelFetch; // Starts fetching a published model on a thread of its own and returns at once. Fetching one -// already here finishes at once. Where it went is waifu_modelmanager_describe_json's to say. +// already here finishes at once. Where it went is waifu_modelmanager_manifest_path's to say. WaifuModelFetch *waifu_modelmanager_fetch_async(const char *model, void *user_data, WaifuProgressCallback on_progress, @@ -246,6 +299,19 @@ void waifu_modelmanager_fetch_cancel(WaifuModelFetch *fetch); // After on_complete; cancels first if it has not had it. void waifu_modelmanager_fetch_free(WaifuModelFetch *fetch); +// -- what a model is ------------------------------------------------------------------------- +// +// libwaifu's, from here on. + +// A JSON document; schema below. `model` is the path of a manifest -- whose weights need not be +// beside it -- or "tones". Free with waifu_string_free. NULL, with waifu_last_error, otherwise. +char *waifu_describe_json(const char *model); + +// Every device, as JSON: {"devices": [{"name", "available", "free"}]}, `free` being the bytes +// free on the card for CUDA and Vulkan where they can be used. Which suits a model is the +// caller's to decide. Asking CUDA builds a context on the card. +char *waifu_devices_json(void); + // -- the engine ------------------------------------------------------------------------------ typedef enum { @@ -274,8 +340,9 @@ void waifu_engine_free(WaifuEngine *engine); void waifu_engine_cancel(WaifuEngine *engine, WaifuJob job); // running: after this step void waifu_engine_cancel_all(WaifuEngine *engine); -// Fetches if needed, then reads onto the device, dropping the model before. What was loaded is -// waifu_modelmanager_describe_json's to say. +// Reads the model whose manifest is at `model` onto the device, dropping the model before; "tones" +// is the stand-in voice. A published model is fetched, and its path found, with the model +// manager first. What was loaded is waifu_describe_json's to say. WaifuJob waifu_engine_load_async(WaifuEngine *engine, const char *model, void *user_data, WaifuProgressCallback on_progress, @@ -414,26 +481,30 @@ change. Everything a job is made of is a struct. ```json { "models": [ { "name": "sdxl:noob", "full_name": "NoobAI XL 1.1", "kind": "image", - "cached": true, "manifest_here": true, "bytes_on_disk": 6938000000, "explicit": false } ] } + "cached": true, "manifest_here": true, "bytes_on_disk": 6938000000, "explicit": false, + "directory": "/Users/me/.cache/libwaifu/models/ling0322--libwaifu-sdxl-noob" } ] } ``` `cached` is every package on the disk; `manifest_here` is its manifest, with or without them -- -what `waifu_modelmanager_describe_json` reads a model's suggestions from. An app can be built with -every manifest (`waifu/examples/fetch_manifests.rs`) and hand them over with +what `waifu_describe_json` reads a model's suggestions from. An app can be built with +every manifest (`cargo run -p waifu-tools --bin fetch_manifests`) and hand them over with `waifu_modelmanager_set_bundled_manifests`, or fetch them with -`waifu_modelmanager_fetch_manifests_async`. +`waifu_modelmanager_fetch_manifests_async`; `waifu_modelmanager_manifest_path` then hands back +one to describe. `directory` is where its packages are, once any of them is: null before. `kind` is `image`, `speech` or `conversion`; a model that does two is listed twice. -`waifu_modelmanager_describe_json` -- the same fields webui's `/api/model` has for its `model`, `voice` or -`converter`, under one object with a `kind`: +`waifu_describe_json` -- what a request for the model may hold, under one object with a `kind`. +`name` is the path it was asked with and `full_name` the manifest's file name; what the catalogue +calls a published one is in the catalogue. `on_disk` is every file the manifest names beside it, +and `bytes` how much they are, null where they are not all there. What to start a screen's boxes +at is not here -- running a model needs none of it -- but in `waifu_modelmanager_defaults_json`, +whose fields a screen lays into these: ```json -{ "kind": "image", "name": "sdxl:noob", "full_name": "NoobAI XL 1.1", - "on_disk": true, "in_memory": false, - "width": 832, "height": 1216, "steps": 28, "guidance": 5.0, "takes_guidance": true, - "sampler": "Euler", "sizes": [[1024, 1024], [832, 1216]], - "prompt": null, "avoid": null, +{ "kind": "image", "name": "/Users/me/.cache/libwaifu/models/.../noob.yaml", + "full_name": "noob.yaml", "on_disk": true, "bytes": 6938000000, "in_memory": false, + "sampler": "Euler", "alignment": 32, "prompt_images": 0, "image_keys": ["start_from"], "why_not": { "prompt_images": "SDXL reads words only", "control": "there is no ControlNet for it" } } @@ -450,6 +521,22 @@ so that a window shows only the boxes that do something: `no_picture_because` was for `start_from` alone. A request that asks for more than this anyway ends with `WAIFU_ERR_MODEL` and the same sentence. +Everything else a request says -- size, steps, guidance, the negative prompt -- is run as given. + +`waifu_modelmanager_defaults_json` -- what a screen starts the boxes at: the kind of model's own +defaults, with the manifest's `suggested:` block laid over them. For a picture model: + +```json +{ "width": 832, "height": 1216, "steps": 28, "guidance": 5.0, "takes_guidance": true, + "sizes": [[832, 1216], [1216, 832], [1024, 1024]], "prompt": "masterpiece", "avoid": null } +``` + +`sizes` is the manifest's own where it names at least three, and the sizes SDXL was trained at +otherwise. `takes_guidance` is false for a kind with no second pass (Krea 2), or where the +manifest says so, or where it suggests a guidance of 1 and says nothing else -- a screen offers +neither a guidance nor a negative prompt then. For a voice it is `{"speed", "temperature"}`, and +for one that also converts `{"speed", "temperature", "conversion": {"steps"}}` -- the +`conversion` going inside `waifu_describe_json`'s. ## Inside the engine @@ -477,22 +564,32 @@ never leave the thread. `Engine` is plain Rust with closures for progress and completion, so cargo test drives it with the `tones` voice, which holds no weights: order, cancelling queued and running jobs, exactly one -`on_complete`, dropping the engine with jobs waiting. `waifu-ffi` wraps it, and its own tests call the +`on_complete`, dropping the engine with jobs waiting. `c_api.rs` wraps it, and its own tests call the `extern "C"` functions the way C would. ## Building -- `waifu-ffi` is a workspace member with `crate-type = ["cdylib"]`, named `waifu`: the library - is `libwaifu.dylib`, `waifu.dll` or `libwaifu.so`. -- Its `build.rs` reads `$LIBWAIFU_LIB_DIR/flint_link_flags.txt` like the crate's, and passes the - link arguments to its own link. On macOS it also sets the install name to - `@rpath/libwaifu.dylib`. -- cbindgen writes `waifu.h` from the `#[repr(C)]` types, so the header cannot drift from the code. -- CMake gets a `waifu-ffi` target beside `waifu-cli`, writing the library and header to `build/`. +- `waifu` and `waifu-modelmanager` are each `crate-type = ["cdylib", "rlib"]`: the library is + what everything outside the crate links, and the rlib is for the crate's own tests and + examples. Cargo names them after their crates -- `libwaifu.dylib`, `libwaifu_modelmanager.dylib` + (`waifu.dll`, `libwaifu.so`, ...) -- and each sets its install name to `@rpath/`. +- Neither crate depends on the other, nor on any crate they share: the two have no code in + common. If one linked the other, every `#[no_mangle]` function of the one it linked would be + exported from both libraries, and a linker could bind an app's call to the wrong copy. What both + need -- the YAML a manifest is written in, the log sink, the C conversions -- each has a copy of. +- libwaifu's `build.rs` reads `$LIBWAIFU_LIB_DIR/flint_link_flags.txt` and passes the link + arguments to its own link. libwaifu_modelmanager links nothing native and builds without CMake. +- cbindgen writes each header from the crate's `src/c_api.rs` and `src/c_api/support.rs`, so the + headers cannot drift from the code. +- CMake's `waifu-engine` and `waifu-modelmanager` targets copy the libraries and their headers + into `build/` and `build/include/` (and on Windows the `.dll.lib` import libraries). + `waifu-cli` then builds the command line, which depends on both crates -- their rlibs -- and so + has them linked in statically: it calls `waifu::c_api` and `waifu_modelmanager::c_api` and + nothing else of theirs, and needs neither library beside it. ## The wrappers -**Swift** (`macos/Sources/WaifuKit`): a module map over `waifu.h`, and +**Swift** (`waifu-app/swiftui/Sources/WaifuKit`): a module map over `waifu.h`, and ```swift final class Engine: Sendable { @@ -531,7 +628,7 @@ job. Both wrappers marshal progress to their UI thread themselves. 2. A log sink in flint, and `log.rs` on the Rust side; with no sink set, every line still prints as it does today. 3. `Engine` in the library, with its tests. -4. `waifu-ffi`, cbindgen, the CMake target, and C-level tests. +4. The C API, cbindgen, the CMake targets, and C-level tests. 5. WaifuKit in Swift, and the app moved onto it. 6. .NET, when it is wanted. diff --git a/docs/indextts.md b/docs/indextts.md index 4a0d4fb..6cebb2e 100644 --- a/docs/indextts.md +++ b/docs/indextts.md @@ -3,8 +3,7 @@ A few seconds of somebody speaking and a sentence in; the sentence, in that voice, out. ```bash -waifu -task text2speech -m models/indextts25.yaml # the text2speech page -cargo run --release --example speak -- models/indextts25.yaml voice.wav "Hello there." out.wav +cargo run --release -p waifu-tools --bin speak -- models/indextts25.yaml voice.wav "Hello there." out.wav ``` ```rust @@ -13,9 +12,10 @@ let voice = tts.listen(&recording)?; // once per voice let sound = tts.say("今天天气很好。", &voice, &options, &mut report)?; ``` -Seven models, each with a module and a document of its own. This one is about the order they run -in, the shapes between them, and the handful of things upstream does in the gaps that no single -module owns. +Seven models run in a row. Six are IndexTTS's alone and are described here: w2v-bert, the GPT, the +emotion path, the semantic codec, S2Mel and the BigVGAN vocoder. [CAMPPlus](campplus.md) is shared +with CosyVoice3 and has a document of its own, as do [the audio front end](audio_features.md) +and [the text normalizer](text_normalize.md). ## Building the package @@ -36,116 +36,268 @@ cd tools && ../.venv/bin/python indextts_package.py -output ../models/indextts25 ``` What comes out is `models/indextts25.yaml`, three weight files and a `tokenizer.json`: 1 367 M -parameters, all of them float32. Nothing is narrowed -- upstream runs S2Mel and the vocoder with -autocast off, and a vocoder is exactly where half a precision's error becomes audible -- so the -package is about the size of the release. +parameters, all float32. Nothing is narrowed: upstream runs S2Mel and the vocoder with autocast +off, and a vocoder is where half a precision's error becomes audible. -Every export was written against the list of tensors its graph actually loads, which -`cargo run --release --example parameters -- ` prints, and verified back against it. +Every exporter keeps the checkpoint's own tensor names and was written against the list the graph +loads, which `cargo run --release -p waifu-tools --bin parameters -- ` prints. Layers that +upstream wraps in `weight_norm` are exported folded, because upstream calls `remove_weight_norm()` +as it loads them. ## The order -**Listening, once per voice.** The recording is brought to 22.05 kHz and cut to fifteen seconds -- -upstream reads it with `librosa.load`, whose default rate that is -- and from there to 16 kHz. +**Listening, once per voice.** The recording is brought to 22.05 kHz and cut to fifteen seconds +-- upstream reads it with `librosa.load`, whose default rate that is -- and from there to 16 kHz. | | reads | gives | | --- | --- | --- | -| [w2v-bert](w2v_bert.md) | the 16 kHz [features](indextts_features.md), 160 wide | `(1, T, 1024)`, standardized | +| [w2v-bert](#w2v-bert-20) | the 16 kHz [features](audio_features.md), 160 wide | `(1, T, 1024)`, standardized | | [CAMPPlus](campplus.md) | Kaldi energies of the 16 kHz audio | who is speaking, `(1, 192)` | -| [the emotion path](indextts_emotion.md) | w2v-bert's features | how they sound, `(1, 1280)` | +| [the emotion path](#the-emotion-path) | w2v-bert's features | how they sound, `(1, 1280)` | | the 22.05 kHz mel | the 22.05 kHz audio | the prompt, `(1, 80, M)` | -| S2Mel's length regulator | w2v-bert's features, stretched onto `M` frames | the prompt's condition, `(1, M, 512)` | +| [S2Mel's length regulator](#s2mel) | w2v-bert's features, stretched onto `M` frames | the prompt's condition, `(1, M, 512)` | **Saying, once per sentence.** The text is cleaned, normalized, prefixed with its language, cut into segments of at most 120 tokens, and each segment goes through: -1. [The GPT](indextts_gpt.md) -- speaker, emotion and text in, semantic tokens out. -2. [The codec's decoder](semantic_codec.md) -- the tokens back to 1024-wide features, two per token. -3. The length regulator -- those stretched onto 1.72 mel frames each, 22 050 / 256 frames a - second over the codec's fifty, divided by the speed. -4. [S2Mel](s2mel.md) -- twenty-five Euler steps with guidance at 0.7, the prompt's mel held in - front and cut off the result afterwards. -5. [BigVGAN](bigvgan.md) -- the mel to 22.05 kHz audio. +1. [The GPT](#the-gpt) -- speaker, emotion and text in, semantic tokens out. +2. [The codec's decoder](#the-semantic-codec) -- the tokens back to 1024-wide features, two per + token. +3. The length regulator -- those stretched onto 1.72 mel frames each (22 050 / 256 frames a second + over the codec's fifty), divided by the speed. +4. [S2Mel](#s2mel) -- twenty-five Euler steps with guidance at 0.7, the prompt's mel held in front + and cut off the result afterwards. +5. [BigVGAN](#bigvgan) -- the mel to 22.05 kHz audio. Segments are joined with 200 milliseconds of silence, as upstream's `interval_silence` does. ## What upstream does in the gaps - **Punctuation is replaced before anything reads the text.** `char_rep_map` turns `，` into `,`, - `。` into `.`, and every bracket and quotation mark into an apostrophe. This is not cosmetic: - without it a Chinese sentence reaches the GPT ending in a `。` it has hardly seen, and it keeps - talking -- Whisper heard "今天天气很好，我们一起去公园散步吧" and then "法尔". With it, three seeds out - of three said the sentence and stopped. The table is applied with upstream's own - first-match-wins semantics, which is why `，，，` becomes three commas and not an ellipsis. + `。` into `.`, and every bracket and quotation mark into an apostrophe. Without it a Chinese + sentence reaches the GPT ending in a `。` it has hardly seen, and it keeps talking. The table is + applied first-match-wins, as upstream does, so `，，，` becomes three commas, not an ellipsis. - **Ids 0 and 1 are taken out of the text.** `prepare_gpt_inputs` drops every start and stop token - before putting one of each back, and those are also ordinary tokens of the vocabulary. So a + before putting one of each back, and those are also ordinary tokens of the vocabulary, so a character that encodes to either is silently not read -- by upstream, and so here. - **The emotion comes from the same recording.** Upstream takes a separate emotion recording when - given one and falls back to the speaker's; `merge_emovec` with the two the same is the speaker's - own emotion vector. + given one and falls back to the speaker's. ## Where this differs from upstream, knowingly -- **No beam search.** Upstream samples inside a three-way beam search. This draws one sequence with - the same temperature, top-k, top-p and repetition penalty. -- **The language is guessed from the script.** Upstream is told it. Kana is Japanese, any other CJK - character is Chinese, everything else English -- wrong for Spanish, which is what - `IndexTts::say_in` is for. The page has no language box yet. -- **The emotion cannot be set.** Upstream can also take eight named emotions as a vector, or read - one out of an instruction with a small Qwen model. Neither is here; the voice sounds the way its - recording felt. +- **No beam search.** Upstream samples inside a three-way beam search; this draws one sequence with + the same temperature, top-k, top-p and repetition penalty. Typical sampling is not offered: + `flint` has no operator for it. +- **The language is guessed from the script.** Kana is Japanese, any other CJK character Chinese, + everything else English -- wrong for Spanish, which is what `IndexTts::say_in` is for. +- **The emotion cannot be set.** Upstream can also take eight named emotions (`emo_matrix`, + `spk_matrix`) or read one out of an instruction with a small Qwen model. The voice here sounds + the way its recording felt. - **w2v-bert's features drop an odd frame at the end** where `SeamlessM4TFeatureExtractor` pads it - and masks it out. Twenty milliseconds, at the edge, masked either way. + and masks it out. Twenty milliseconds, masked either way. +- **No padding masks anywhere.** Upstream pads batches and carries length masks through w2v-bert, + the emotion path and S2Mel. One recording and one segment at a time makes every position valid. -## What it took from `flint` +## The models + +### w2v-bert-2.0 + +Meta's conformer, which reads speech into the features everything else conditions on. + +- **Sixteen layers of twenty-four.** IndexTTS reads `hidden_states[17]` -- the projection plus + sixteen layers -- so the last eight are neither exported nor run. `Config::USED_LAYERS` is that + sixteen, and stopping early has a test of its own. +- **Macaron.** The two feed forwards are added back at half weight, `x * 0.5 + residual`. +- **The depthwise convolution is causal**, padding `kernel - 1` on the left and nothing on the + right -- unlike every other depthwise convolution in this repository. +- **Positions are relative keys**, clamped to sixty-four back and eight forward: seventy-three + embeddings. The index table is built on the host and `lookup` turns it into `(T, T, head_dim)`, + which is the memory to watch. + +### The GPT + +A stock `GPT2Model` -- twenty-four layers, 1280 wide, twenty heads -- fed embeddings this module +assembles rather than ids. + +``` +[ zero padding ][ speaker + emotion ][ 0 ][ 0 ][ text tokens ][ start mel ] +``` -Running in float32 on a card found three things the half-precision models never had: +- **Its position embedding is replaced by zeros.** Upstream deletes `gpt.wpe`; positions come from + learned tables applied to the text and to the mel separately. Leaving GPT-2's own `wpe` in gives + a model wrong in every layer and right in every shape. +- **`Conv1D` is not `Linear`.** HuggingFace's GPT-2 stores its weights `(in, out)` and is read + without a transpose; `mel_head` is an `nn.Linear` and is transposed. +- **The activation is `gelu_new`**, the tanh approximation, not the exact GELU. They differ by a + thousandth per layer, which does not survive twenty-four. +- **The cache is ordinary graph values, not paged attention.** `paged_attention` and + `store_kv_cache` are eager operators and need the FlashAttention kernels, so a stack built on + them could be neither one compiled graph nor tested on the CPU. A step concatenates its keys onto + the history, and `flint`'s causal mask is aligned to the bottom right, so one block serves both + `Gpt::prefill` and `Gpt::step`. +- **The head scores the last position only**, taken before the eight-thousand-wide projection. +- **The start token is at position zero**, so the first generated token is at position one. Off by + one here still sounds like speech, which is what makes it easy to miss. + +The fast tests compare against GPT-2 imported from the pinned transformers, every width and the +layer count cut, in three probes: the stack, the head, the embeddings. The loop is checked by an +invariant instead -- a cached step says what the whole sequence read in one pass says -- and two +lessons came out of that: compare **scores, not tokens** (an argmax survived every mutation +tried), and draw the weights **scattered** (xorshift), because swept weights make every key point +the same way and the resulting flat softmax does not notice the history being concatenated +backwards. + +The released weights have a test too, because weights made up from their own names agree with any +self-consistent transposition: `lang_embedding` was declared with nine rows against a table of 107 +(`len(LANGUAGE_DICT) + 1`), every fast test passed, and only the checkpoint said otherwise. The +exporter writes the GPT's 301 tensors and the emotion path's 153 out of the same `gpt.pth`; only +`text_head` is left behind. + +### The emotion path + +A recording in, the one 1280-wide row that sits beside the speaker's in the GPT's prefix. + +| stage | what it is | out | +| --- | --- | --- | +| `emo_conditioning_encoder` | a WeNet conformer: a convolution that halves the length, four blocks | `(T/2, 512)` | +| `emo_perceiver_encoder` | lucidrains' `PerceiverResampler` with **one** latent, so a vector | `(1, 1024)` | +| `emovec_layer`, `emo_layer` | two projections | `(1280)` | + +- **One matrix holds 134 M parameters.** `embed.out.0.weight` is `(512, 261632)`: WeNet's + `Conv2dSubsampling2` convolves the 1024 features as an image and flattens 512 channels by 511 + columns into one projection. The exporter checks the shape, since it is the tensor a reader is + most likely to "fix". +- **The relative attention has no shift.** WeNet's `rel_shift` is commented out in the released + source, so the position term is an ordinary matrix multiply. +- **The gated feed forward gates the second half**: `gelu(second) * first`, the other way round + from `flint`'s `geglu`, so the two slices are written out. +- **It is not macaron**, where w2v-bert is: one feed forward at full weight. +- **The position table is the checkpoint's, rounded through bfloat16.** The released `pos_enc.pe` + equals the sinusoid cast to `bfloat16`, bit for bit, and upstream runs with it because loading + overwrites the exact table. So the exporter writes the buffer out (ten megabytes) and checks that + equality exactly. + +The tests run the released `ConformerEncoder` and `PerceiverResampler`, in five probes cut where +upstream cuts. The weights are drawn at 0.7 rather than the usual 0.25: smaller, and GELU is so +nearly linear that swapping the gate's halves went unnoticed. + +### The semantic codec + +MaskGCT's codec, whose 8192 entries are the alphabet the GPT writes. + +- **The half that runs is the decoder.** `infer_v2_5.py` decodes the GPT's tokens into features + for the length regulator; the line that would quantize the reference recording is commented out. + So the exporter leaves the encoder out, though the module keeps it, as the definition of what a + token means. +- **The decoder is a second backbone**, structurally the encoder's, not its transpose. Two + asymmetries are the release's: `quantize` has a GELU after `down` and `decode` none after + `out_project`; and the doubling is a nearest interpolation, written as the sequence concatenated + with itself on a new last axis and flattened. +- **The nearest entry is found on the host.** `flint` has no argmax, so the graph computes the + similarity -- both sides L2 normalized, so nearest is the largest dot product -- and `nearest` + walks it. Its own test covers an all-negative row and a tie, which the earlier index wins, as in + `torch.max`. + +The reference harness writes a two-function torchaudio shim into the download cache rather than +installing torchaudio, which would move the venv's pinned torch. + +### S2Mel + +A flow-matching model: given a partly denoised mel and a time it returns a direction, and +`solve_euler` integrates it. It is conditioned on the semantic tokens (`cond`), the prompt's mel +(`prompt_x`, zeros past it -- what makes it continue a voice), CAMPPlus's `style`, and the time. + +A thirteen-layer transformer (RMS norm, rotary positions, SwiGLU, adaptive norms, U-net skips from +half the layers to the other half), then an eight-layer WaveNet. + +- **The rotary embedding is interleaved**, pairing `2i` with `2i + 1`, not `flint`'s NeoX `i` with + `i + D/2`; `rotate` does it out of cosine and sine tables, as Anima does. +- **The content is continuous whatever the configuration says.** `content_type: 'discrete'` + builds a `cond_embedder`, but `forward` always uses `cond_projection`. +- **The final modulation is `(1 + scale) * norm(x) + shift`**, where every other adaptive norm is + `weight * norm(x) + bias`. +- **The sampler holds the prompt**: its frames are forced back after every step. **Guidance is two + passes**, `(1 + rate) * conditioned - rate * unconditioned`, rather than upstream's batch of two. +- **The length regulator resizes the way torch does**, with the scale computed and multiplied in + float32, because an exact division picks a different frame for about one length in seventy. +- **Mish is written without a logarithm**: `tanh(softplus(x)) = (u² + 2u) / (u² + 2u + 2)` with + `u = exp(x)`, exact, and safe because every call follows a group norm. + +The tests keep every count and cut every width -- thirteen and eight layers, since the skip +connections are a function of depth -- and check the trunk, the denoiser and the regulator against +upstream, plus four closed forms: Mish, a constant velocity integrating to exactly one unit, +guidance as two passes, and no guidance as one. + +### BigVGAN + +`nvidia/bigvgan_v2_22khz_80band_256x`: 112 M parameters, 80 mel bands in, 22.05 kHz out, 256 +samples per frame. NVIDIA's, MIT licensed; it sits under `indextts` because nothing else here runs +it, and moves beside [`waifu::audio`](audio.md) the day something does. + +It adds no operator: six transposed convolutions, each followed by three averaged residual blocks, +and 109 anti-aliased snakes, all of them `audio` functions built from `conv2d`, `matmul` and +arithmetic -- so it runs on every backend. + +- **The snake is `x + sin(alpha x)^2 / beta`**, applied at twice the rate and filtered back down + so the harmonics it makes do not fold into the audible band. +- **The filter is a twelve-tap Kaiser-windowed sinc** (cutoff a quarter, transition 0.3), not the + Hann-windowed `resample_kernel`; the two are not interchangeable. +- **Padding is by replication**, and the upsampling's asymmetric trims are the reference's: get one + wrong and the waveform shifts by a sample. +- **`groups = C` with the same taps everywhere** is read as `(N * C, 1, L)` and one ungrouped + convolution. +- **`alpha` and `beta` are stored as logarithms** and raised where read, as upstream does. + +The activation is where the vocoder's time goes: about six passes over the tensor where a fused +kernel would read it once. NVIDIA ships a CUDA kernel for exactly this, which is the measurement. +`resblock: "2"` is refused rather than run as the other kind; no v2 release uses it. -- **`lookup` failed to launch past 65 535 ids.** The ids sat on the grid's y and z axes, which hold - that many blocks each; w2v-bert asks for one per *pair* of frames, a quarter of a million for ten - seconds. The grid is one-dimensional now, one thread per output element. -- **Reductions were half-only, and rank-three at most.** `sum` over a float32 tensor, or over any - tensor of rank four, stopped with "not implemented". The kernel was always generic over the type; - the dispatch was not. It reads any rank as `(rows, last)` now, with the rows on the x axis for - the same reason as `lookup`. -- **The repetition penalty never ran on a card at all.** It was written for half logits, and the - GPT casts its logits to float before penalizing them; it allowed a history of at most sixty-four - tokens; and a token said twice was two threads racing on one score. It is a gather and then a - scatter now, as the host's is, so a repeated token is penalized once. +## What it took from `flint` -Each has a CUDA test of its own under `flint/cuda/`. +Running in float32 on a card found three things the half-precision models never had, each with a +CUDA test of its own under `flint/cuda/`: + +- **`lookup` failed to launch past 65 535 ids**, which w2v-bert reaches in ten seconds. The grid is + one-dimensional now. +- **Reductions were half-only and rank three at most.** They read any rank as `(rows, last)` now. +- **The repetition penalty never ran on a card**: half logits only, a history of at most + sixty-four, and a repeated token was two threads racing. It is a gather and a scatter now. ## What is checked +In the fast suite, each against upstream's own code -- imported or downloaded and run, not +transcribed -- with weights filled from a hash of each parameter's name on both sides: + | | | | --- | --- | -| `waifu/tests/indextts_features.rs` | the front end, resampling and the 22.05 kHz mel against upstream's own code | -| `indextts::tests` | segment packing, punctuation against upstream's own regular expression, language ids | -| `waifu/tests/indextts.rs` | the whole pipeline on real weights -- ignored, see below | +| `waifu/tests/w2v_bert.rs` | the projection, two layers and four | +| `waifu/tests/indextts_gpt.rs` | the stack, the head, the embeddings, and the cached loop | +| `waifu/tests/indextts_emotion.rs` | the sinusoid table and four cuts through the path | +| `waifu/tests/semantic_codec.rs` | the similarity, and `nearest` | +| `waifu/tests/s2mel.rs` | trunk, denoiser, regulator, and four closed forms | +| `waifu/tests/bigvgan.rs` | a small model with every kind of layer the release has | +| `waifu/tests/indextts_features.rs` | the front end and resampling | +| `indextts::tests` | segment packing, punctuation, language ids | + +The reference scripts are `tools/_reference.py`, which print the Rust constants. -The whole-pipeline test needs the package and two recordings: +Ignored, needing the package or a checkpoint: the GPT and BigVGAN against their released weights, +and the whole pipeline, which needs two recordings: ```bash PYTHONPATH=~/.cache/libwaifu/pydeps .venv/bin/python tools/indextts_reference_voices.py -cargo test --release --manifest-path waifu/Cargo.toml --test indextts -- --ignored --nocapture +cargo test --release --manifest-path waifu/Cargo.toml --no-fail-fast \ + --test indextts --test indextts_gpt -- --ignored --test-threads=1 ``` -It asks, in one pass: that a sentence comes out as a sound of a sentence's length; that it sounds -more like the speaker it was given than like another one, by CAMPPlus's own embedding; that one -seed says it identically twice; that a stop comes back with nothing; and that asking with no -recording is a sentence saying why. - -What no test here can check is whether the words are right, since Whisper is not a thing this crate -runs. That was checked by hand while this was written, and the numbers are worth keeping: - -- **Whisper heard every sentence exactly**: an English one, a Chinese one on three seeds, and a - 24-second passage that crossed several segments -- numbers included, "25 dollars and 99 cents" - coming back as "$25.99". One Chinese sentence came back with 界面 heard as the near-homophone 见面. -- **It is the voice it was given.** Upstream's CAMPPlus, over three speakers from LibriSpeech: - 0.85, 0.85 and 0.89 between each output and its own recording, 0.32 to 0.63 against the others. - Two real recordings of one speaker sit at 0.84. -- **A seed repeats**, across processes, byte for byte. - -On an RTX 5060 Ti: the package reads in about a second and a half, a recording is listened to in -under a second, and a four-second sentence takes four seconds -- about one for the GPT and three for -S2Mel's fifty passes and the vocoder. +The pipeline test asks that a sentence comes out a sentence long, sounds more like its own speaker +than another by CAMPPlus's embedding, repeats exactly on one seed, stops cleanly, and refuses with +a reason when given no recording. Whether the words are right needs Whisper, which was checked by +hand: it heard every sentence exactly, English and Chinese, numbers included ("25 dollars and 99 +cents" as "$25.99"), except 界面 once heard as 见面. Upstream's CAMPPlus put each output at 0.85 +to 0.89 from its own speaker and 0.32 to 0.63 from the others; two real recordings of one speaker +sit at 0.84. + +On an RTX 5060 Ti the package reads in about a second and a half, a recording is listened to in +under a second, and a four-second sentence takes four seconds: about one for the GPT and three for +S2Mel and the vocoder. diff --git a/docs/indextts_emotion.md b/docs/indextts_emotion.md deleted file mode 100644 index 2dc2edd..0000000 --- a/docs/indextts_emotion.md +++ /dev/null @@ -1,158 +0,0 @@ -# IndexTTS-2.5's emotion path - -Where the emotion vector comes from when nobody supplies one. A recording in, one row of 1280 out -— the row that sits beside the speaker's in [the GPT](indextts_gpt.md)'s prefix. - -```rust -use waifu::indextts::emotion::{self, Config}; - -let config = Config::indextts(); -let emotion = indextts_emotion::graph( - &g, features, &config, frames, positions, dtype, device, -)?; -``` - -`features` is what [w2v-bert](w2v_bert.md) read off the reference audio — the same tensor -[the semantic codec](semantic_codec.md) is handed, standardized by the release's own mean and -deviation. `positions` is the sinusoid table, and there is a wrinkle about which copy of it to -use: see [below](#the-position-table-is-not-the-closed-form). - -## Four stages - -| stage | what it is | out | -| --- | --- | --- | -| `emo_conditioning_encoder` | a WeNet conformer: one convolution that halves the length, four blocks | `(T/2, 512)` | -| `emo_perceiver_encoder` | lucidrains' `PerceiverResampler`, with **one** latent | `(1, 1024)` | -| `emovec_layer` | a projection | `(1280)` | -| `emo_layer` | another | `(1280)` | - -153 tensors and 163 M parameters, all under the checkpoint's own names — this module renames -nothing and folds nothing, for the reason [the GPT](indextts_gpt.md) gives. - -The one latent is the whole trick of the middle stage. The speaker path's perceiver resamples a -recording to thirty-two latents and gets a sequence; this one asks for a single latent, and what -comes back is therefore a vector. A feeling is one thing about a recording, so it is one row. - -## 134 M parameters in a single matrix - -`emo_conditioning_encoder.embed.out.0.weight` is `(512, 261632)`. That is more than the conformer -it belongs to, and a fifth of everything the GPT's package holds. - -It is not a transcription error. WeNet's `Conv2dSubsampling2` reads the 1024 features as an image -one channel deep, runs a 3×3 convolution at stride two over *both* axes, and then flattens -everything it produced at each output frame — 512 channels by the 511 feature columns that -survive 1024 — into one projection down to 512. Five hundred and eleven columns times five hundred -and twelve channels is the 261,632. - -Worth knowing before anyone wonders why a four-block conformer weighs as much as a good part of -the twenty-four-layer stack it feeds, and worth checking at export, which -`tools/indextts_gpt_exporter.py` does: it is the one tensor here whose size a reader is most -likely to assume is wrong and helpfully "fix". - -## Three things upstream does that are easy to get wrong - -**The relative attention has no shift.** The score is Transformer-XL's — `(q + u) · k` plus -`(q + v) · p`, where `p` is the sinusoid table through `linear_pos` — but WeNet's `rel_shift`, -which is what would turn the second term's absolute positions into relative ones, is commented -out in the released source with a note that it is useless for speech. So the position term is an -ordinary matrix multiply, and a reimplementation that helpfully puts the shift back is wrong. - -**The gated feed forward gates the second half.** `flint`'s `geglu` computes `gelu(first) * -second`; this `GEGLU` is `gelu(second) * first`. The halves are the other way round, so the module -writes the two slices out rather than calling the operator, and the weight stays laid out the way -the checkpoint holds it. Both use the exact GELU, which is the one thing about the two that does -agree. - -**This conformer is not macaron and [w2v-bert's](w2v_bert.md) is.** There, two feed forwards are -added back at half weight; here there is one, at full. The two conformers in this one pipeline -differ on exactly that, which is how the mistake would arrive — and a scale copied from the wrong -one is wrong in every layer without being wrong in any shape. - -## There is no mask - -Upstream pads a batch of recordings out to the longest and carries a length mask through both -stages; `emo_cond_mask_pad` is the `ConstantPad1d` that prepends a `True` for the one latent. One -recording at a time — which is what a reading is — makes every position valid, so the masked -softmax is an ordinary one and there is nothing to pad. - -That also makes the perceiver's context order irrelevant: concatenating the latents after the -frames rather than before permutes the keys and the values together, and a softmax does not care. -It would start to matter the moment a batch were padded, which is why the mask exists upstream and -why leaving it out is a statement about the caller rather than about the arithmetic. - -## The position table is not the closed form - -`emo_conditioning_encoder.embed.pos_enc.pe` is a registered buffer — five thousand rows of -sinusoid that depend on nothing but their own shape — so the obvious thing is to rebuild it at -load and save the ten megabytes. That turns out to be wrong, and it is the one surprise this -module turned up. - -**The released table differs from the sinusoid at essentially every entry**, by 7.5e-4 on average -and 2.0e-3 at worst. Chasing that number down is worth the five minutes it takes, because the -first guess is wrong and the second one settles the design: - -| hypothesis | verdict | -| --- | --- | -| every stored value is representable in `float16` | **no** | -| every stored value is representable in `bfloat16` | **yes** | -| `stored == sinusoid(...).to(float16)` | no, 2.0e-3 out | -| `stored == sinusoid(...).to(bfloat16)` | **yes, difference exactly zero** | - -So it is not a mystery and not a corruption: somebody saved this model through `bfloat16` once, -and the buffer kept the scar. The worst case lands where it should: 2.0e-3 is 2^-9, exactly half -a `bfloat16` step for a value just under one, which is the most round-to-nearest can be out with a -7-bit significand. - -And upstream runs with the rounded one. `pe` is persistent, so it is in the file; -`load_checkpoint` calls `load_state_dict(strict=False)`, which only forgives *absent* keys and -this one is present; so the exact table computed in `PositionalEncoding.__init__` is overwritten -the moment the checkpoint loads. A reimplementation that rebuilds the table from the formula is -therefore running a model the release does not run — by very little, but for no reason. - -That leaves a genuine choice, and it could be made either way. The exporter writes the tensor out: -ten megabytes is 0.3% of the package, against a standing claim that a Rust loop over `f64` sines -lands on the same `bfloat16` grid `torch` reached through `float32`. True as far as it was tested, -and not worth having to keep true. The check at export is now a bit-for-bit equality rather than a -tolerance — a mystery that has been solved deserves a check that fails the moment it comes back. - -`waifu::indextts::emotion::positional_encoding` still builds the unrounded table, for the tests and -for running without a package. What that costs is 0.13% of the position term after `linear_pos`, -below the noise floor of the half precision this model runs in on a card — worth writing down -rather than worth worrying about. - -## What is checked - -`waifu/tests/indextts_emotion.rs`, against the released source's own `ConformerEncoder` and -`PerceiverResampler` — imported and run, not transcribed. That distinction earns its keep here: -every one of the mistakes listed above produces a tensor of the right shape, and none would be -caught by comparing this against a second reading of the paper. The reference is -`tools/indextts_emotion_reference.py`, which needs no download and no checkpoint. - -```bash -cargo test --manifest-path waifu/Cargo.toml --test indextts_emotion -``` - -Five probes: the sinusoid table, which both sides build for themselves and which fails before any -weight is involved; then the model cut where upstream cuts it, after `self.embed`, after the -blocks, after the perceiver, and after the two projections. A failure says which. - -Nine plausible bugs were written into the module one at a time and every one of them fails, by -between 55 and 11,000 times the tolerance. **Getting there needed the weights drawn larger than -the other reference scripts here draw them.** At the usual 0.25 the activations land where GELU is -still nearly a straight line, `gelu(a) * b` and `gelu(b) * a` agree to a thousandth, and the -end-to-end probe did not notice the two halves being swapped at all. Drawing at 0.7 puts the -activations at the order of one, which is where the released model runs and where its -nonlinearities are nonlinear. - -They are also drawn scattered rather than swept, for the reason `docs/indextts_gpt.md` sets out at -length: a swept parameter makes every key in a projection point nearly the same way, and the flat -softmax that follows is an average, which does not care what order its keys arrived in. Attention -is most of what this module is. - -## What this does not do - -`emo_matrix`, `spk_matrix` and the `qwen0.6bemo4-merge` text model in the release are the -*other* ways 2.5 can be told a feeling: eight named emotions with vectors to interpolate between, -and a small language model that reads an instruction like "say it sadly" and produces a mixture. -This is the path that takes the feeling off the reference recording, which is what -`inference_speech` does when it is given neither. diff --git a/docs/indextts_gpt.md b/docs/indextts_gpt.md deleted file mode 100644 index 098680f..0000000 --- a/docs/indextts_gpt.md +++ /dev/null @@ -1,158 +0,0 @@ -# IndexTTS-2.5's GPT - -Text and a voice in, the semantic tokens S2Mel reads out. The largest of the models here and the -one the rest of the pipeline is arranged around — and what it generates is neither audio nor mel, -but the codec's alphabet, one token per two frames of the speech to come. - -## It is an ordinary GPT-2, fed unusually - -Twenty-four layers, 1280 wide, twenty heads: a stock `GPT2Model`. Two things are taken out of it. - -Its **token embedding is unused** — what goes in is embeddings this module assembles, not ids. - -Its **position embedding is replaced by zeros**. The reference deletes `gpt.wpe` and puts -`null_position_embeddings` there. The positions this model uses are added earlier, by a learned -table applied to the text and to the mel separately. A reader who assumes GPT-2's own `wpe` is in -play gets a model that is wrong in every layer and still returns a tensor of the right shape, -which is why the backbone has a test of its own. - -## What the prefix is made of - -``` -[ zero padding ][ speaker + emotion ][ 0 ][ 0 ][ text tokens ][ start mel ] -``` - -Three conditioning rows — one carrying the speaker vector from [CAMPPlus](campplus.md) plus an -emotion vector, then two rows of zeros that are not learned and carry nothing — and then the -text, each token embedded and given both a position and a language. The padding is on the -**left**, so generation always begins the same distance from the end. - -## Two shapes that are not what they look like - -**`Conv1D` is not `Linear`.** HuggingFace's GPT-2 uses its own `Conv1D`, whose weight is -`(in, out)` where `nn.Linear` stores `(out, in)`. `conv1d_linear` reads it as stored and does not -transpose. Using the ordinary `Linear` layer here would load the right numbers in the wrong order -and produce plausible nonsense. Note that `mel_head` *is* an `nn.Linear` and does transpose — the -two conventions sit ten lines apart. - -**The activation is the tanh approximation.** `gelu_new`, not the exact GELU `flint` has. The two -differ by about a thousandth, which is small enough to look like noise in one layer and not small -enough to survive twenty-four. `gelu_new` writes it out: `0.5x(1 + tanh(√(2/π)(x + 0.044715x³)))`. - -## What is checked - -`waifu/tests/indextts_gpt.rs`, in the fast suite. **No download at all** — `GPT2Model` is in the -transformers this venv already pins, so `tools/indextts_gpt_reference.py` imports the reference -rather than fetching it, and zeroes `wpe` to match what the reference does by deleting it. - -Every width is cut and so is the layer count, because a GPT-2 stack has no bookkeeping that turns -on its depth. Three probes kept apart so a failure says which part: the stack, the head on top, -and the embeddings in front. - -```bash -cargo test --manifest-path waifu/Cargo.toml --test indextts_gpt -.venv/bin/python tools/indextts_gpt_reference.py -``` - -## The generation loop - -`Gpt` is the model with weights behind it and the loop on top: two graphs, compiled once and run -many times. - -| | reads | keeps | -| --- | --- | --- | -| `Gpt::prefill` | conditioning, text, start token — the whole prefix at once | every layer's keys and values | -| `Gpt::step` | one token | the same, one position longer | - -`Gpt::generate` is the loop over the two, with `Sampling` deciding what to draw and when to stop. -Both are public, so a caller that wants a sampler this crate does not offer can run them itself. - -### Why it is not paged attention - -An earlier draft of this page said the cache should be built on `paged_attention` and -`store_kv_cache`, after the `KVCacheManager` pattern in the author's `libllm`. It is not, for two -reasons. - -Those two are **eager** operators — `op.rs` says so in as many words, because an operation whose -result is a change to somebody else's tensor has no way to say that in a graph. A stack built on -them could not be one compiled `Ir`. And `paged_attention` rides on the FlashAttention kernels, -so a build without them has no paged attention at all and no portable fallback: the fast suite, -which is where this is tested, runs on the CPU. - -What is here instead is the cache as ordinary graph values. A step concatenates its own keys onto -the history and attends over the lot, and `flint`'s causal mask is **aligned to the bottom right -of the score matrix** — so one query against a longer history sees all of it with `causal` left -true, and the same block serves the prefill and the step with no second path written. The -concatenation copies the history once per step, which is the same order of memory traffic the -attention beside it was always going to spend reading it. - -What the cache saves is the rest: twenty-four layers of projections and a head eight thousand -wide, over the whole prefix, at every one of up to 1815 steps. - -### Two places it is easy to be wrong - -The head scores the **last position only**, taken before the matmul rather than after. A prefill -over a six-hundred-token sentence would otherwise compute six hundred rows of an eight-thousand- -wide projection and sample one of them. - -The start token is a mel token at **position zero**, so the first token generated is at position -one. A loop that counts its own output from zero puts every token one place from where the model -was trained to find it — and that is the quiet one, because the reading still sounds like speech. - -### What the loop is checked against - -The three probes above compare against GPT-2. The loop cannot: there is no reference constant for -a reading. What a cached step has to satisfy instead is that it says what the whole sequence says, -read in one pass — an invariant that needs no second implementation to state. - -Two things that cost time and are worth knowing before touching those tests: - -- **The scores, not the tokens.** A token is an argmax, and an argmax over a forty-token alphabet - survives a great deal of being wrong. Every mutation this was written against — the history - concatenated on the wrong side, a mel position off by one, the prefix assembled in the wrong - order — left the argmax alone and moved the scores. -- **The weights have to scatter.** `Filled` sweeps one sine per parameter, so every row of a - projection samples the same curve, every key it produces points nearly the same way, and the - softmax over them is flat. A flat softmax is an average, and an average does not care what order - its keys arrived in: the whole cache test passed with the history concatenated backwards. The - loop tests draw their weights by xorshift instead. Nothing there is compared against Python, so - the only thing the numbers have to be is different from one another. - -### And against the released weights - -`tools/indextts_gpt_exporter.py` writes the 301 tensors this graph reads out of the release's -`gpt.pth` — and, since it is the same checkpoint, the 153 more that -[the emotion path](indextts_emotion.md) reads. It renames and folds nothing: both modules were -written against the checkpoint's own names, so a tensor that turns out to be wrong can be compared -with the one it came from. Of the 456 in the file, only `text_head` is left behind. - -```bash -.venv/bin/python tools/indextts_gpt_exporter.py \ - -checkpoint ~/.cache/libwaifu/indextts25/gpt.pth \ - -output models/indextts25-gpt.safetensors -.venv/bin/python tools/indextts_gpt_real_reference.py - -cargo test --release --manifest-path waifu/Cargo.toml \ - --test indextts_gpt -- --ignored --test-threads=1 -``` - -That test is worth having for one reason the fast ones cannot cover: **weights made up from their -own names agree with any transposition of themselves that is self-consistent**, and the `Filled` -source answers every shape it is asked for. `lang_embedding` is the case that actually happened — -it was declared with nine rows against a table that has a hundred and seven, every fast test -passed, and only the real checkpoint said otherwise. 107 is `len(LANGUAGE_DICT) + 1`: Whisper's -list of 106 languages, most of which this model was never trained to say. - -## Where `emo_vec` comes from - -`prefill` takes the finished 1280-wide emotion vector directly, which is the path -`inference_speech` takes when it is given one. When it is not, the four stages that make one out -of the reference recording are [`waifu::indextts::emotion`](indextts_emotion.md) — -`emo_conditioning_encoder` → `emo_perceiver_encoder` → `emovec_layer` → `emo_layer`, 153 of the -checkpoint's tensors, exported by the same `tools/indextts_gpt_exporter.py` that writes this -model. - -## What is still missing - -- **Typical sampling**, which the reference offers and `flint` has no operator for. It is left out - rather than approximated with one that is nearby. diff --git a/docs/s2mel.md b/docs/s2mel.md deleted file mode 100644 index 1250d8a..0000000 --- a/docs/s2mel.md +++ /dev/null @@ -1,148 +0,0 @@ -# S2Mel - -The middle of IndexTTS-2.5. The GPT in front of it produces semantic tokens; this turns them into -the mel spectrogram a BigVGAN makes audio of. - -```rust -use waifu::indextts::s2mel::{self, Config}; - -let direction = s2mel::graph( - &g, x, prompt_x, cond, style, time, time2, - &Config::indextts(), frames, cos, sin, dtype, device, -)?; -``` - -It is a **flow-matching** model, so the network does not produce a mel. It produces one step of a -trajectory — given a partly denoised mel and a time, the direction to move — and a sampler calls -it a handful of times. The network, the sampler that drives it and the length regulator in -front of it are all here. - -## What it is conditioned on - -Four different things, which is worth spelling out because they arrive separately: - -| | | -| --- | --- | -| `cond` | the semantic tokens, one 512-wide vector per output frame | -| `prompt_x` | the reference speaker's own mel where there is one and zeros after, which is what makes this *continue* a voice rather than invent one | -| `style` | 192 numbers from [`campplus`](campplus.md) — who is speaking | -| `t` | where along the trajectory this step is | - -## The shape of it - -A thirteen-layer transformer over the frame axis, then an eight-layer WaveNet over the same axis. - -The transformer is `gpt_fast`'s rather than a diffusion model's: RMS normalization, rotary -positions, a SwiGLU feed forward. What makes it a *diffusion* transformer is that every -normalization is adaptive — the timestep projects to a weight and a bias that scale and shift it -— and that half the layers hand a skip connection forward to the other half, U-net fashion. - -## Three things that are not what a reader would guess - -### The rotary embedding is interleaved, not split - -`flint`'s `rotary_embedding` rotates the NeoX way, pairing element `i` with element `i + D/2`. -This model pairs `2i` with `2i + 1` — the other convention entirely, and it gives different -numbers for the same weights. So `rotate` does it explicitly, out of cosine and sine tables the -caller passes in, the same way `anima`'s transformer does. - -The tables are `(T, head_dim / 2)` and the tensor is laid out with heads before time, so they -line up against the last two axes and broadcast over the rest without anything being moved. - -### The content is read as continuous, whatever the configuration says - -`config.yaml` sets `content_type: 'discrete'` and the model builds a `cond_embedder` for it. But -the reference's `forward` assigns `cond_in_module = self.cond_projection` unconditionally, with -the line that would have chosen between them commented out just above it. So the embedding is -never reached and the projection always is. This follows the code, not the configuration. - -### The final modulation is not the adaptive norm - -Every adaptive normalization in the transformer applies the projected weight as it comes: -`weight * norm(x) + bias`. The `FinalLayer` at the end does `(1 + scale) * norm(x) + shift` -instead, over a `LayerNorm` with no affine of its own. Two different conventions in one model. - -## What is checked - -`waifu/tests/s2mel.rs`, in the fast suite, against IndexTTS's own implementation — -`tools/s2mel_reference.py` downloads the reference and runs it. No package and no download at -test time: every parameter is filled from its own name, by the same FNV-1a hash and the same sine -in both languages. - -The model under test has **every width cut and every count kept** — 16 mel bands, a 64-wide -transformer, four heads, but still thirteen transformer layers and eight WaveNet layers. Depth is -what the skip connections turn on: half the layers hand one forward to the other half, taken off -the end of a list, so which layer receives which is a function of the depth and nothing else. At a -depth of two there is nothing to get wrong; at thirteen there is an off-by-one in three places. - -Seven tests. Three compare against the reference, so a failure lands on one side or the other; -four check arithmetic that has a closed form and needs no reference at all. - -| test | what it covers | -| --- | --- | -| `the_transformer_trunk_is_the_reference_trunk` | the skips and the rotary convention | -| `the_denoiser_is_the_reference_denoiser` | end to end, including the WaveNet and the final modulation | -| `the_length_regulator_is_the_reference_regulator` | nine tokens onto twenty-four frames, a ratio that is not whole | -| `mish_is_the_activation_it_replaces` | the exp-only identity, against the definition | -| `the_solver_integrates_a_constant_velocity_exactly` | one unit of time, however many steps it is cut into | -| `guidance_is_two_passes_combined` | the combination, and which prompt each pass is handed | -| `no_guidance_is_one_pass` | that a rate of zero costs one pass and not two | - -```bash -cargo test --manifest-path waifu/Cargo.toml --test s2mel -.venv/bin/python tools/s2mel_reference.py # to regenerate the constants -``` - -## The sampler - -`solve_euler` is the loop around the denoiser: `steps` even steps from 0 to 1, moving the mel at -each by `dt` times the direction the network returns. Three things make it more than that loop. - -**The prompt is held, not predicted.** The reference speaker's own mel is laid into a zero tensor -the length of the whole utterance, handed to the network every step, and the frames it covers are -forced back to zero in `x` after every step — so the model never denoises the part it was given. - -**Guidance is two passes.** With `cfg_rate` above zero the network runs again with its -conditioning removed and the two combine as `(1 + rate) * conditioned - rate * unconditioned`. -The caller does the removing, since what has to be zeroed lives in the graph it built; the -closure is told which pass it is being asked for. The reference stacks the two into a batch of -two instead — the same arithmetic, since attention does not mix batch entries, and running twice -keeps the graph at the batch of one everything else here assumes. - -**The direction is not the answer.** What comes back is a velocity; the mel is what integrating -it produces. - -It is tested against an integral that can be done by hand — a constant velocity field over -`steps` steps of `1 / steps` adds up to exactly one unit of time — rather than against a recorded -number, so the test says the loop is right and not merely unchanged. - -## The length regulator - -One vector per semantic token in, one per mel frame out: a projection, a nearest-neighbour resize -onto the mel's frame rate, and four convolution-normalization-activation stages. - -**The resize is torch's, float32 and all.** Upstream's `F.interpolate(size=frames, -mode="nearest")` has output frame `i` copy input frame `min(floor(i * scale), from - 1)`, with -`scale = from / to` worked out *and multiplied* in float32. Where `i * from / to` is a whole -number the float32 product can land a hair under it and floor one frame early — 90 tokens onto -154 frames copies frame 44 at position 77, where the exact quotient says 45 — so an exact integer -division picks a different frame for about one length in seventy. `upsample_nearest1d` computes -the index torch's kernel does, on every backend. - -**Mish has no logarithm in it.** `x * tanh(softplus(x))` needs `ln`, which `flint` does not have. -With `u = exp(x)` the identity `tanh(softplus(x)) = (u² + 2u) / (u² + 2u + 2)` is exact and uses -only exp, multiply, add and divide. What is not exact is `exp` of a large number, which overflows -a float32 above about 88; every call here is on the far side of a group normalization, so the -numbers are a few standard deviations and nowhere near it. The test checks the identity against -the definition over a range either side of zero. - -## What is still missing - -**The exporter.** There is none yet; `tools/s2mel_reference.py` fetches the reference source but -never the 415 MB checkpoint. Note for whoever writes it: several layers are wrapped in -`weight_norm`, which the reference removes when it loads, so the folded weight is what to export -— the same thing `tools/bigvgan_exporter.py` does. - -**Padding masks.** The reference builds a mask from `x_lens` and passes it to the attention. For -a batch of one at full length that mask is all true, which is what this assumes. A batch of -several utterances of different lengths would need it. diff --git a/docs/semantic_codec.md b/docs/semantic_codec.md deleted file mode 100644 index d72e1a0..0000000 --- a/docs/semantic_codec.md +++ /dev/null @@ -1,100 +0,0 @@ -# The semantic codec - -The alphabet the GPT reads and writes, and the way back out of it. MaskGCT's codec, as -IndexTTS-2.5 vendors it. - -```rust -// What the pipeline runs: tokens back into features. -let features = semantic_codec::decode(&g, codes, &Config::indextts(), tokens, dtype, device)?; - -// What defines the alphabet, and what this release never calls. -let scores = semantic_codec::similarity(&g, features, &Config::indextts(), frames, dtype, device)?; -let tokens = semantic_codec::nearest(&scores, config.codebook_size as usize); -``` - -## The half that runs is the decoder, which is backwards from the guess - -A text-to-speech pipeline sounds like it should encode. It does not. `infer_v2_5.py` calls - -```python -S_infer = self.semantic_codec.decode(codes) -``` - -to turn the GPT's tokens into the features S2Mel's length regulator reads — and the line that -would run the encoder over the reference audio, - -```python -# _, S_ref = self.semantic_codec.quantize(spk_cond_emb) -``` - -is **commented out in the released source**, replaced by `S_ref = self.get_emb(...)`. The -reference recording reaches the length regulator as w2v-bert's continuous 1024-wide features, -never having been near a codebook. - -So `decode` is what a package carries and what a reading runs. The encoder stays in this module -because it is the definition of the alphabet — the thing that says what a token *means* — but -nothing here needs it to say a sentence, and `tools/semantic_codec_exporter.py` leaves it out. - -An earlier version of this document had it the other way round. It was wrong. - -## The shape of it - -Encoding: a stride-two convolution halves the frame rate, then twelve ConvNeXt blocks read what is -left, then every frame is projected down to eight numbers and replaced by the nearest of 8192 -codebook entries. The token is that entry's index. - -Decoding is the mirror, at the same shapes: the entry is looked up, `out_project` widens those -eight numbers back to 1024, **a second backbone with its own weights** — structurally identical to -the encoder's, and not its transpose, because a codec is not an invertible function — reads them, -and then the frame rate is doubled back. One token becomes two frames, which is the ratio it stood -for. - -Two asymmetries worth knowing, both of them the release's rather than this port's. `quantize` puts -a GELU after `down`; `decode` puts none after `out_project`. And the doubling is -`F.interpolate(scale_factor=2, mode="nearest")` followed by a width-three convolution — a repeat -*is* nearest interpolation at exactly two, so the graph concatenates the sequence with itself -along a new last axis and flattens it, which puts every frame beside its own copy. - -A ConvNeXt block is a depthwise convolution seven frames wide that mixes no channels, then two -linears that mix channels and no frames. The depthwise one is `conv1d` with a group per channel. - -## The nearest entry is found on the host - -`flint` can take a maximum but cannot say *where* it was: there is no argmax. So the graph -computes the similarity — `(T, 8) @ (8, 8192)` — and `nearest` walks the result on the host. - -That is less of a compromise than it sounds. Both sides are L2 normalized, which makes the -Euclidean distance `2 - 2 * cosine`, so the nearest entry is the largest dot product and the -whole search is one matrix multiply the graph is good at. What comes back is `T * 8192` floats -once per utterance, and walking them is short beside the twelve ConvNeXt blocks in front. - -## What is checked - -`waifu/tests/semantic_codec.rs`, against the reference — `tools/semantic_codec_reference.py` -downloads it and runs it. Twelve ConvNeXt layers as shipped, every width cut: a 32-wide feature -instead of 1024, a 64-entry codebook instead of 8192. - -The search gets its own test rather than resting on the codes this small model picks. With -sixty-four entries and made-up weights it picks the same one for every frame, which would pass -against almost any indexing mistake. So the *similarity* is what checks the encoder, and -`nearest` is checked against a matrix whose answer is obvious by construction — including a row -that is negative throughout, which catches an implementation that starts its search from zero, -and a tie, which the earlier index has to win because that is what `torch.max` does. - -```bash -cargo test --manifest-path waifu/Cargo.toml --test semantic_codec -.venv/bin/python tools/semantic_codec_reference.py -``` - -## A note on the reference harness - -`vocos.py` imports two mel helpers from torchaudio at module scope, for a class `VocosBackbone` -never touches. torchaudio is not installed in this venv and should not be — it pins itself -against a torch version, and this venv's pins are what hold several models' reference tensors -steady. So `_torchaudio_shim` writes the two functions into the download cache instead, with -their real arithmetic rather than a stub that raises: a shim that lied would be worse than none. - -## What is missing - -The exporter. `codec.pth` is 607 MB and nothing here fetches it yet. The decoder half, which -training needs and inference does not. diff --git a/docs/speech.md b/docs/speech.md index 37461ed..7206c07 100644 --- a/docs/speech.md +++ b/docs/speech.md @@ -3,7 +3,7 @@ > **The page is gone.** The web page this was written around was removed on 2026-10-13, when > drawing went back to the terminal. Speech now runs through [`waifu::engine`](../waifu/src/engine.rs) > and the C API over it ([ffi.md](ffi.md)), which the macOS app speaks with; `waifu` in the -> terminal draws only, and `examples/speak.rs` says a sentence from a checkout. What follows about +> terminal draws only, and waifu-tools' `speak` says a sentence from a checkout. What follows about > the [`Voice`](../waifu/src/speech.rs) trait, the models and the WAV reader still holds; what it > says about the page, its tab and `-m` for a voice is history. @@ -176,8 +176,5 @@ shape and neither has the other. Nothing here needs a model, a card or a network, so all of it runs in the fast suite: ```bash -LIBWAIFU_LIB_DIR="$PWD/build" cargo test --manifest-path waifu/Cargo.toml --features cli --lib +LIBWAIFU_LIB_DIR="$PWD/build" cargo test --manifest-path waifu/Cargo.toml --lib ``` - -`--features cli` matters: without it, none of `waifu/src/cli` is compiled, so none of the routes, -the state or the worker is tested. The library half -- `speech` and `wav` -- is tested either way. diff --git a/docs/indextts_normalize.md b/docs/text_normalize.md similarity index 96% rename from docs/indextts_normalize.md rename to docs/text_normalize.md index 89f64ba..fea7a65 100644 --- a/docs/indextts_normalize.md +++ b/docs/text_normalize.md @@ -1,11 +1,12 @@ -# IndexTTS-2.5's text normalizer +# The text normalizer Writing out what a number is read as, so the model never sees a digit. `$5.20` becomes "five point two zero dollars", `共1234人` becomes `共一千二百三十四人`, and `Son 21 días` becomes `Son veintiuno días`. -`waifu::indextts::normalize`, beside [the GPT](indextts_gpt.md) rather than at the crate root — -a diffusion model has no text to normalize, and this is the speech pipeline's. +`waifu::text_normalizer`, written for [IndexTTS-2.5](indextts.md) and used by +[CosyVoice3](cosyvoice3.md) too -- a diffusion model has no text to normalize, and this is the +speech pipelines'. ## Three languages, and one that is skipped diff --git a/docs/w2v_bert.md b/docs/w2v_bert.md deleted file mode 100644 index 92f19d9..0000000 --- a/docs/w2v_bert.md +++ /dev/null @@ -1,67 +0,0 @@ -# w2v-bert-2.0 - -The conformer that reads speech into the features everything else conditions on. Meta's model, -shared by several speech systems; IndexTTS-2.5 runs it over the reference audio and hands the -result to [the semantic codec](semantic_codec.md). - -## Sixteen layers of twenty-four - -IndexTTS reads `hidden_states[17]` — the projection plus sixteen layers. **The last eight are -never used.** That is a third of a 580 M parameter model that does not have to run, or be -exported, at inference. `Config::USED_LAYERS` is that sixteen, and `graph` takes the count. - -Stopping partway has its own test, because "the stack is right" and "reading it early lands in -the right place" are different claims and only the second is what `hidden_states[17]` rests on. - -## A conformer layer is four things, and two of them are halved - -Feed forward, attention, convolution, feed forward, then a normalization. The two feed forwards -are added back at **half weight** — `x * 0.5 + residual`, which is what "macaron" means. It is -not a detail: at full weight the residual stream doubles through every layer. - -## Two things worth pointing at - -### The depthwise convolution is causal - -Every other depthwise convolution in this repository pads symmetrically. This one pads -`kernel - 1` on the left and nothing on the right, so a frame never sees the future. -`conv1d` takes one padding for both sides, so the padding happens in front of it and -the convolution is asked for none. - -### The positions are relative keys, not rotary - -`position_embeddings_type` is `relative_key`: the score between two frames gets a term that -depends only on how far apart they are, clamped to sixty-four back and eight forward — so there -are seventy-three embeddings and the rest of the distance is thrown away. - -Which embedding applies to which pair is a table of indices built on the host, and `lookup` turns -it into the `(T, T, head_dim)` the scores need. That makes the gather an operator that exists. -The term itself is one matrix multiply per query frame, with the frame as the batch axis: -`(T, B*H, D)` against `(T, D, T)`. - -The `(T, T, head_dim)` is the memory to watch. It is the same for every layer and could be built -once and shared if that ever matters; at present each layer builds its own. - -## What is checked - -`waifu/tests/w2v_bert.rs`, in the fast suite. **No download of source or of weights** — -`Wav2Vec2BertModel` is in the pinned transformers already in the venv, so -`tools/w2v_bert_reference.py` imports the reference rather than fetching it. - -Every width is cut, and so is the layer count — four rather than twenty-four — because a -conformer stack has no bookkeeping that turns on its depth: layer *n* reads layer *n-1* and that -is all. What is kept is the shape of one layer, which is where this model is unusual. - -Three probes: the projection alone, two layers, and four. - -```bash -cargo test --manifest-path waifu/Cargo.toml --test w2v_bert -.venv/bin/python tools/w2v_bert_reference.py -``` - -## What is missing - -The exporter, and the feature extractor in front of the model — `SeamlessM4TFeatureExtractor`, -which turns audio into the 160-wide stacked mel features this reads. Padding masks: the reference -masks padded frames out of both the attention and the depthwise convolution, and a batch of one -at full length needs neither. diff --git a/macos/Sources/CWaifu/include/waifu.h b/macos/Sources/CWaifu/include/waifu.h deleted file mode 120000 index 53e8e88..0000000 --- a/macos/Sources/CWaifu/include/waifu.h +++ /dev/null @@ -1 +0,0 @@ -../../../../waifu-ffi/include/waifu.h \ No newline at end of file diff --git a/tools/indextts_emotion_reference.py b/tools/indextts_emotion_reference.py index 5a3b18a..1873101 100644 --- a/tools/indextts_emotion_reference.py +++ b/tools/indextts_emotion_reference.py @@ -44,7 +44,7 @@ The weights are drawn by xorshift from each parameter's own name, in both languages, and *not* by sweeping a sine the way the other reference scripts here do. The reason is written up in -`docs/indextts_gpt.md`: a swept parameter makes every key in a projection point nearly the same +`docs/indextts.md`, under the GPT: a swept parameter makes every key in a projection point nearly the same way, the softmax over them comes out flat, and a flat softmax is an average -- which does not care what order its keys arrived in. Attention is most of what this module is, so it is drawn scattered or it is not being tested. diff --git a/tools/indextts_gpt_exporter.py b/tools/indextts_gpt_exporter.py index e86e4d8..d088efd 100644 --- a/tools/indextts_gpt_exporter.py +++ b/tools/indextts_gpt_exporter.py @@ -333,7 +333,7 @@ def report_position_buffer(state): f"{POSITION_BUFFER} is no longer the sinusoid rounded to bfloat16 -- {wrong} of " f"{off_rounded.numel()} elements are more than one bfloat16 step ({coarsest_bfloat16_step:.3g}) " f"from the closed form, worst {off.max().item():.3g}. The release has changed its " - f"position encoding, and docs/indextts_emotion.md is now telling a story about the old one" + f"position encoding, and docs/indextts.md is now telling a story about the old one" ) return positions, dim, off.mean().item(), off.max().item() diff --git a/tools/indextts_normalize_reference.py b/tools/indextts_normalize_reference.py index d22f556..83693b6 100644 --- a/tools/indextts_normalize_reference.py +++ b/tools/indextts_normalize_reference.py @@ -97,7 +97,7 @@ def ours(cases): extension of it to normalize a string would be a great deal of machinery for a pure function. """ program = """ -use waifu::indextts_normalize::{normalize, Language}; +use waifu::text_normalizer::{normalize, Language}; fn main() { let mut input = String::new(); diff --git a/tools/semantic_codec_exporter.py b/tools/semantic_codec_exporter.py index 540ff69..ae24bce 100644 --- a/tools/semantic_codec_exporter.py +++ b/tools/semantic_codec_exporter.py @@ -66,7 +66,7 @@ That the 121 names the graph loads are all present and all the right shape, and that folding produced the shape the graph declares. - cargo run --release --example parameters -- semantic_codec_decode + cargo run --release -p waifu-tools --bin parameters -- semantic_codec_decode prints that list, and it is the authority; this file is written against it. diff --git a/tools/w2v_bert_exporter.py b/tools/w2v_bert_exporter.py index b451158..06c5c01 100644 --- a/tools/w2v_bert_exporter.py +++ b/tools/w2v_bert_exporter.py @@ -34,7 +34,7 @@ `waifu::w2v_bert::Config::USED_LAYERS` is that sixteen, and this file's `LAYERS` has to agree with it. They are two constants in two languages saying the same thing, so the count is printed at the -end for a human to compare against `cargo run --example parameters -- w2v_bert`. +end for a human to compare against `cargo run -p waifu-tools --bin parameters -- w2v_bert`. # Nothing is done to the weights diff --git a/macos/.gitignore b/waifu-app/swiftui/.gitignore similarity index 100% rename from macos/.gitignore rename to waifu-app/swiftui/.gitignore diff --git a/macos/Package.swift b/waifu-app/swiftui/Package.swift similarity index 61% rename from macos/Package.swift rename to waifu-app/swiftui/Package.swift index 116e72c..75e663f 100644 --- a/macos/Package.swift +++ b/waifu-app/swiftui/Package.swift @@ -1,16 +1,19 @@ // swift-tools-version: 5.10 -// The macOS app: native controls over libwaifu, which it links. WaifuKit is the Swift face of the -// C API in waifu-ffi (docs/ffi.md); CWaifu is that API's header, as a module Swift can import. +// The macOS app: native controls over libwaifu and libwaifu_modelmanager, which it links. WaifuKit +// is the Swift face of the C APIs in waifu and waifu-modelmanager (docs/ffi.md); CWaifu is their +// headers, as a module Swift can import. // -// libwaifu.dylib is CMake's to build, into build/ at the top of the repository: this links it -// from there, and finds it there at run time while the app runs unbundled. Inside libwaifu-gui.app it is -// in Contents/Frameworks, which the second rpath is for. +// Both dylibs are CMake's to build, into build/ at the top of the repository: this links them from +// there, and finds them there at run time while the app runs unbundled. Inside libwaifu-gui.app +// they are in Contents/Frameworks, which the second rpath is for. import Foundation import PackageDescription -let repository = URL(fileURLWithPath: #filePath).deletingLastPathComponent().deletingLastPathComponent() +// This file is waifu-app/swiftui/Package.swift: the repository is three levels up from it. +let repository = URL(fileURLWithPath: #filePath) + .deletingLastPathComponent().deletingLastPathComponent().deletingLastPathComponent() let build = repository.appendingPathComponent("build").path let package = Package( @@ -22,6 +25,7 @@ let package = Package( path: "Sources/CWaifu", linkerSettings: [ .linkedLibrary("waifu"), + .linkedLibrary("waifu_modelmanager"), .unsafeFlags(["-L", build]), ] ), diff --git a/macos/README.md b/waifu-app/swiftui/README.md similarity index 69% rename from macos/README.md rename to waifu-app/swiftui/README.md index 76389a7..65f2192 100644 --- a/macos/README.md +++ b/waifu-app/swiftui/README.md @@ -3,28 +3,29 @@ A native SwiftUI app on libwaifu. It needs macOS 14. The app does none of the drawing itself, and runs no `waifu` program: it calls the library in -process, through `libwaifu.dylib` and its C API (`docs/ffi.md`), wrapped for Swift in `WaifuKit` -(`Sources/WaifuKit`). +process, through `libwaifu.dylib` and `libwaifu_modelmanager.dylib` and their C APIs +(`docs/ffi.md`), wrapped for Swift in `WaifuKit` (`Sources/WaifuKit`). ## Build ```bash -# the library, from the top of the repository: build/libwaifu.dylib and build/include/waifu.h, -# for macOS 14 and later (docs/metal.md) +# the libraries, from the top of the repository: build/libwaifu.dylib, +# build/libwaifu_modelmanager.dylib and their headers in build/include, for macOS 14 and later +# (docs/metal.md) cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DWITH_METAL=ON cmake --build build -j$(sysctl -n hw.ncpu) -# the app, with build/libwaifu.dylib in its Frameworks, every published model's manifest in its +# the app, with both libraries in its Frameworks, every published model's manifest in its # Resources (fetched now, a few KB each), and its icon compiled in -macos/build-app.sh -open macos/build/libwaifu-gui.app +waifu-app/swiftui/build-app.sh +open waifu-app/swiftui/build/libwaifu-gui.app -# the disk image, macos/build/libwaifu-gui.dmg: the app beside an Applications link, at 128 pixels -macos/package-dmg.sh +# the disk image, waifu-app/swiftui/build/libwaifu-gui.dmg: the app beside an Applications link, at 128 pixels +waifu-app/swiftui/package-dmg.sh ``` `package-dmg.sh` lays the window out with `dmg_layout.py`, which writes the image's `.DS_Store` -with the `ds_store` package; the script installs it the first time into `macos/build/dmg-venv`. +with the `ds_store` package; the script installs it the first time into `waifu-app/swiftui/build/dmg-venv`. The release workflow runs both scripts and attaches it to the release as `libwaifu-gui--macos-arm64.dmg`. @@ -36,16 +37,17 @@ both with a Developer ID and have Apple notarize the image: xcrun notarytool store-credentials waifu-notary --apple-id --team-id export WAIFU_SIGN_IDENTITY="Developer ID Application: ()" -macos/build-app.sh -WAIFU_NOTARY_PROFILE=waifu-notary macos/package-dmg.sh +waifu-app/swiftui/build-app.sh +WAIFU_NOTARY_PROFILE=waifu-notary waifu-app/swiftui/package-dmg.sh ``` The release workflow does the same from the secrets `MACOS_CERT_P12` (the certificate and its key, exported as .p12, in base64), `MACOS_CERT_PASSWORD`, `MACOS_SIGN_IDENTITY`, and an App Store Connect API key: `ASC_KEY_P8` (in base64), `ASC_KEY_ID` and `ASC_ISSUER_ID`. -While working on it, `swift run --package-path macos Waifu` runs it without a bundle, against -`build/libwaifu.dylib`; `swift test --package-path macos` checks WaifuKit against it. +While working on it, `swift run --package-path waifu-app/swiftui Waifu` runs it without a bundle, against +`build/libwaifu.dylib` and `build/libwaifu_modelmanager.dylib`; `swift test --package-path waifu-app/swiftui` +checks WaifuKit against them. The icon is `Resources/AppIcon.icon`, the Icon Composer format, which `build-app.sh` compiles with `actool` (Xcode 26 or later). Its one layer is `Assets/picture.png`, 1024 pixels square. diff --git a/macos/Resources/AppIcon.icon/Assets/picture.png b/waifu-app/swiftui/Resources/AppIcon.icon/Assets/picture.png similarity index 100% rename from macos/Resources/AppIcon.icon/Assets/picture.png rename to waifu-app/swiftui/Resources/AppIcon.icon/Assets/picture.png diff --git a/macos/Resources/AppIcon.icon/icon.json b/waifu-app/swiftui/Resources/AppIcon.icon/icon.json similarity index 100% rename from macos/Resources/AppIcon.icon/icon.json rename to waifu-app/swiftui/Resources/AppIcon.icon/icon.json diff --git a/waifu-app/swiftui/Sources/CWaifu/include/waifu.h b/waifu-app/swiftui/Sources/CWaifu/include/waifu.h new file mode 120000 index 0000000..2eca728 --- /dev/null +++ b/waifu-app/swiftui/Sources/CWaifu/include/waifu.h @@ -0,0 +1 @@ +../../../../../waifu/include/waifu.h \ No newline at end of file diff --git a/waifu-app/swiftui/Sources/CWaifu/include/waifu_modelmanager.h b/waifu-app/swiftui/Sources/CWaifu/include/waifu_modelmanager.h new file mode 120000 index 0000000..0185ccc --- /dev/null +++ b/waifu-app/swiftui/Sources/CWaifu/include/waifu_modelmanager.h @@ -0,0 +1 @@ +../../../../../waifu-modelmanager/include/waifu_modelmanager.h \ No newline at end of file diff --git a/macos/Sources/CWaifu/shim.c b/waifu-app/swiftui/Sources/CWaifu/shim.c similarity index 79% rename from macos/Sources/CWaifu/shim.c rename to waifu-app/swiftui/Sources/CWaifu/shim.c index 2b7d84d..321cf8f 100644 --- a/macos/Sources/CWaifu/shim.c +++ b/waifu-app/swiftui/Sources/CWaifu/shim.c @@ -17,6 +17,8 @@ // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -// The C API is libwaifu.dylib's; this target only makes waifu.h a module Swift can import. A C -// target needs one source file to be a target at all, and this is it. +// The C API is libwaifu.dylib's and libwaifu_modelmanager.dylib's; this target only makes their +// headers a module Swift can import. Each header stands alone and the two share no type. A C target +// needs one source file to be a target at all, and this is it. #include "waifu.h" +#include "waifu_modelmanager.h" diff --git a/macos/Sources/Waifu/Launcher.swift b/waifu-app/swiftui/Sources/Waifu/Launcher.swift similarity index 100% rename from macos/Sources/Waifu/Launcher.swift rename to waifu-app/swiftui/Sources/Waifu/Launcher.swift diff --git a/macos/Sources/Waifu/MainView.swift b/waifu-app/swiftui/Sources/Waifu/MainView.swift similarity index 100% rename from macos/Sources/Waifu/MainView.swift rename to waifu-app/swiftui/Sources/Waifu/MainView.swift diff --git a/macos/Sources/Waifu/Media.swift b/waifu-app/swiftui/Sources/Waifu/Media.swift similarity index 100% rename from macos/Sources/Waifu/Media.swift rename to waifu-app/swiftui/Sources/Waifu/Media.swift diff --git a/macos/Sources/Waifu/ModelCard.swift b/waifu-app/swiftui/Sources/Waifu/ModelCard.swift similarity index 100% rename from macos/Sources/Waifu/ModelCard.swift rename to waifu-app/swiftui/Sources/Waifu/ModelCard.swift diff --git a/macos/Sources/Waifu/ModelManagerView.swift b/waifu-app/swiftui/Sources/Waifu/ModelManagerView.swift similarity index 100% rename from macos/Sources/Waifu/ModelManagerView.swift rename to waifu-app/swiftui/Sources/Waifu/ModelManagerView.swift diff --git a/macos/Sources/Waifu/Models.swift b/waifu-app/swiftui/Sources/Waifu/Models.swift similarity index 99% rename from macos/Sources/Waifu/Models.swift rename to waifu-app/swiftui/Sources/Waifu/Models.swift index e992c0e..ac9e46f 100644 --- a/macos/Sources/Waifu/Models.swift +++ b/waifu-app/swiftui/Sources/Waifu/Models.swift @@ -31,7 +31,7 @@ private let decoder: JSONDecoder = { // MARK: - a model -/// A model that draws, as `waifu_modelmanager_describe_json` says it. +/// A model that draws, as `waifu_describe_json` says it. struct Chosen: Decodable { var name: String var fullName: String? diff --git a/macos/Sources/Waifu/OutputPane.swift b/waifu-app/swiftui/Sources/Waifu/OutputPane.swift similarity index 100% rename from macos/Sources/Waifu/OutputPane.swift rename to waifu-app/swiftui/Sources/Waifu/OutputPane.swift diff --git a/macos/Sources/Waifu/Session.swift b/waifu-app/swiftui/Sources/Waifu/Session.swift similarity index 100% rename from macos/Sources/Waifu/Session.swift rename to waifu-app/swiftui/Sources/Waifu/Session.swift diff --git a/macos/Sources/Waifu/SettingsColumn.swift b/waifu-app/swiftui/Sources/Waifu/SettingsColumn.swift similarity index 100% rename from macos/Sources/Waifu/SettingsColumn.swift rename to waifu-app/swiftui/Sources/Waifu/SettingsColumn.swift diff --git a/macos/Sources/Waifu/WaifuApp.swift b/waifu-app/swiftui/Sources/Waifu/WaifuApp.swift similarity index 100% rename from macos/Sources/Waifu/WaifuApp.swift rename to waifu-app/swiftui/Sources/Waifu/WaifuApp.swift diff --git a/macos/Sources/WaifuKit/WaifuKit.swift b/waifu-app/swiftui/Sources/WaifuKit/WaifuKit.swift similarity index 70% rename from macos/Sources/WaifuKit/WaifuKit.swift rename to waifu-app/swiftui/Sources/WaifuKit/WaifuKit.swift index 3f22a88..fbad7c4 100644 --- a/macos/Sources/WaifuKit/WaifuKit.swift +++ b/waifu-app/swiftui/Sources/WaifuKit/WaifuKit.swift @@ -59,19 +59,37 @@ public enum WaifuError: Error, LocalizedError, Equatable { default: self = .internalError(message) } } + + /// The same, from libwaifu_modelmanager, which has codes of its own. + init(code: WaifuModelManagerStatusCode, message: String) { + switch code { + case WAIFU_MODELMANAGER_ERR_INVALID_ARGUMENT: self = .refused(message) + case WAIFU_MODELMANAGER_ERR_UNKNOWN_MODEL: self = .unknownModel(message) + case WAIFU_MODELMANAGER_ERR_FETCH: self = .fetch(message) + default: self = .internalError(message) + } + } } -/// This thread's last message from the library. +/// This thread's last message from libwaifu. func lastError() -> String { guard let said = waifu_last_error() else { return "" } return String(cString: said) } -/// The library's ABI, and the one this was compiled against: a mismatch is a libwaifu.dylib from -/// another build than waifu.h. +/// This thread's last message from libwaifu_modelmanager, which keeps its own. +func modelManagerLastError() -> String { + guard let said = waifu_modelmanager_last_error() else { return "" } + return String(cString: said) +} + +/// The libraries' ABIs, and the ones this was compiled against: a mismatch is a library from +/// another build than its header. public enum ABI { public static var library: UInt32 { waifu_abi_version() } public static var header: UInt32 { UInt32(WAIFU_ABI_VERSION) } + public static var modelManagerLibrary: UInt32 { waifu_modelmanager_abi_version() } + public static var modelManagerHeader: UInt32 { UInt32(WAIFU_MODELMANAGER_ABI_VERSION) } } // MARK: - what goes in and comes out @@ -185,7 +203,6 @@ public enum Stage: Sendable { init(_ stage: WaifuStage) { switch stage { - case WAIFU_STAGE_FETCHING: self = .fetching case WAIFU_STAGE_READING: self = .reading case WAIFU_STAGE_ENCODING: self = .encoding case WAIFU_STAGE_DRAWING: self = .drawing @@ -222,6 +239,19 @@ public struct JobProgress: Sendable { words = event.words.map { String(cString: $0) } ?? "" seconds = event.seconds } + + /// A download's progress, as libwaifu_modelmanager says it. + init(fetching event: WaifuModelManagerProgressEvent) { + stage = .fetching + fraction = event.fraction < 0 ? nil : event.fraction + done = event.done + total = event.total + part = Int(event.part) + parts = Int(event.parts) + file = event.file.map { String(cString: $0) } + words = event.words.map { String(cString: $0) } ?? "" + seconds = event.seconds + } } public typealias ProgressHandler = @Sendable (JobProgress) -> Void @@ -281,6 +311,28 @@ private let onComplete: WaifuCompleteCallback = { userData, event in pending.finish(event.pointee.status.code, message(event.pointee.status), ()) } +/// libwaifu_modelmanager's progress and ending, which are its own types, ended as a call is. +private let onFetchProgress: WaifuModelManagerProgressCallback = { userData, event in + guard let userData, let event else { return } + let pending = Unmanaged.fromOpaque(userData).takeUnretainedValue() + pending.progress?(JobProgress(fetching: event.pointee)) +} + +private let onFetchComplete: WaifuModelManagerCompleteCallback = { userData, event in + guard let userData, let event else { return } + let pending = Unmanaged.fromOpaque(userData).takeRetainedValue() + let status = event.pointee.status + let said = status.message.map { String(cString: $0) } + switch status.code { + case WAIFU_MODELMANAGER_OK: pending.finish(WAIFU_OK, nil, ()) + case WAIFU_MODELMANAGER_CANCELLED: pending.finish(WAIFU_CANCELLED, nil, nil) + case WAIFU_MODELMANAGER_ERR_UNKNOWN_MODEL: pending.finish(WAIFU_ERR_UNKNOWN_MODEL, said, nil) + case WAIFU_MODELMANAGER_ERR_FETCH: pending.finish(WAIFU_ERR_FETCH, said, nil) + case WAIFU_MODELMANAGER_ERR_INVALID_ARGUMENT: pending.finish(WAIFU_ERR_INVALID_ARGUMENT, said, nil) + default: pending.finish(WAIFU_ERR_INTERNAL, said, nil) + } +} + private let onImage: WaifuImageCompleteCallback = { userData, event in guard let userData, let event else { return } let pending = Unmanaged.fromOpaque(userData).takeRetainedValue() @@ -350,6 +402,7 @@ final class Scratch { func call( progress: ProgressHandler?, cancel: @escaping (WaifuJob) -> Void, + refusal: () -> String = lastError, submit: (UnsafeMutableRawPointer) -> WaifuJob ) async throws -> T? { let ticket = Ticket() @@ -368,7 +421,7 @@ func call( // Refused at the door: nothing will ever be called, so the retain is this side's // to take back. Unmanaged.fromOpaque(userData).release() - continuation.resume(throwing: WaifuError.refused(lastError())) + continuation.resume(throwing: WaifuError.refused(refusal())) return } ticket.given(job, cancel: cancel) @@ -400,12 +453,14 @@ public final class Engine: @unchecked Sendable { { [handle] job in waifu_engine_cancel(handle, job) } } - /// Fetches `model` if it has to, then reads it onto the device, dropping the one before. - /// Throws CancellationError where it was cancelled. + /// Fetches `model` if it has to -- a published name, through libwaifu_modelmanager -- then + /// reads it onto the device, dropping the one before. `model` is a published name, the path + /// of a manifest, or "tones". Throws CancellationError where it was cancelled. public func load(_ model: String, progress: ProgressHandler? = nil) async throws { + let manifest = try await ModelManager.fetchedManifest(model, progress: progress) let loaded: Void? = try await call(progress: progress, cancel: cancel) { userData in - model.withCString { model in - waifu_engine_load_async(handle, model, userData, onProgress, onComplete) + manifest.withCString { manifest in + waifu_engine_load_async(handle, manifest, userData, onProgress, onComplete) } } if loaded == nil { throw CancellationError() } @@ -484,28 +539,119 @@ public final class Engine: @unchecked Sendable { // MARK: - the model manager -/// What is published, what is here, and fetching it. None of it needs an engine. +/// What is published, what is here, and fetching it: libwaifu_modelmanager's. None of it needs an +/// engine; describeJSON asks libwaifu what the manifest the model manager found says. public enum ModelManager { - private static func json(_ said: UnsafeMutablePointer?) throws -> Data { - guard let said else { throw WaifuError.refused(lastError()) } - defer { waifu_string_free(said) } - return Data(String(cString: said).utf8) - } - /// The published models, as the library's JSON: `{"models": [...]}`. public static func catalogJSON() throws -> Data { - try json(waifu_modelmanager_catalog_json()) + guard let said = waifu_modelmanager_catalog_json() else { + throw WaifuError.refused(modelManagerLastError()) + } + defer { waifu_modelmanager_string_free(said) } + return Data(String(cString: said).utf8) } - /// What a model is without reading it, as the library's JSON. Throws for a name that is + /// What a model is without reading it, as the library's JSON: read by libwaifu out of the + /// manifest libwaifu_modelmanager finds for it, whose weights need not be here. Throws for a + /// name with no manifest here -- fetch it, or every manifest, first -- and for one that is /// neither published nor a manifest on the disk. public static func describeJSON(_ model: String) throws -> Data { - try json(model.withCString { waifu_modelmanager_describe_json($0) }) + let manifest = model == Self.tones ? model : try manifestPath(model) + guard let said = manifest.withCString({ waifu_describe_json($0) }) else { + throw WaifuError.refused(lastError()) + } + defer { waifu_string_free(said) } + let described = try withDefaults(Data(String(cString: said).utf8), manifest) + return try named(described, model) + } + + /// What libwaifu says a model can be asked for, with what libwaifu_modelmanager says its + /// boxes start at laid in: the two halves of one description. A `conversion` is in both, and + /// its defaults go inside it. + static func withDefaults(_ described: Data, _ manifest: String) throws -> Data { + guard let said = manifest.withCString({ waifu_modelmanager_defaults_json($0) }) else { + throw WaifuError.refused(modelManagerLastError()) + } + defer { waifu_modelmanager_string_free(said) } + guard var whole = try JSONSerialization.jsonObject(with: described) as? [String: Any], + let defaults = try JSONSerialization.jsonObject(with: Data(String(cString: said).utf8)) + as? [String: Any] + else { return described } + for (key, value) in defaults { + if var inner = whole[key] as? [String: Any], let more = value as? [String: Any] { + inner.merge(more) { _, new in new } + whole[key] = inner + } else { + whole[key] = value + } + } + return try JSONSerialization.data(withJSONObject: whole) + } + + /// libwaifu calls a model by its manifest's path and file name; one asked for by its published + /// name is called that, and what the catalogue calls it, here too. + static func named(_ described: Data, _ model: String) throws -> Data { + guard let listed = try? listed(model), + var said = try JSONSerialization.jsonObject(with: described) as? [String: Any] + else { return described } + said["name"] = model + said["full_name"] = listed.full_name + if var conversion = said["conversion"] as? [String: Any] { + conversion["name"] = model + conversion["full_name"] = listed.full_name + said["conversion"] = conversion + } + return try JSONSerialization.data(withJSONObject: said) + } + + /// The stand-in voice, which is in libwaifu itself and has no manifest. + static let tones = "tones" + + /// The manifest of a published model here -- fetched, or bundled -- or a manifest's own path, + /// handed back as it is: what libwaifu is given. Throws where there is none here. + public static func manifestPath(_ model: String) throws -> String { + guard let said = model.withCString({ waifu_modelmanager_manifest_path($0) }) else { + throw WaifuError.unknownModel(modelManagerLastError()) + } + defer { waifu_modelmanager_string_free(said) } + return String(cString: said) + } + + /// The manifest of `model` with its weights beside it, fetching it first where it is a + /// published model that is not all here. What a load is handed. + static func fetchedManifest(_ model: String, progress: ProgressHandler?) async throws -> String { + if model == tones || FileManager.default.fileExists(atPath: model) { return model } + if try !isCached(model) { try await fetch(model, progress: progress) } + return try manifestPath(model) + } + + /// A published model, as the catalogue lists it. + struct Listed: Decodable { + var name: String + var full_name: String + var cached: Bool + } + + /// What the catalogue says of `model`. Throws for a name it does not list. + static func listed(_ model: String) throws -> Listed { + struct Catalog: Decodable { + var models: [Listed] + } + let catalog = try JSONDecoder().decode(Catalog.self, from: catalogJSON()) + guard let listed = catalog.models.first(where: { $0.name == model }) else { + throw WaifuError.unknownModel("there is no model called \"\(model)\"") + } + return listed + } + + /// Whether every package of a published model is here, as the catalogue says. + static func isCached(_ model: String) throws -> Bool { + try listed(model).cached } public static var directory: String? { guard let said = waifu_modelmanager_directory() else { return nil } - defer { waifu_string_free(said) } + defer { waifu_modelmanager_string_free(said) } return String(cString: said) } @@ -513,7 +659,7 @@ public enum ModelManager { public static func setDirectory(_ path: String?) throws { let code = path.map { $0.withCString { waifu_modelmanager_set_directory($0) } } ?? waifu_modelmanager_set_directory(nil) - if code != WAIFU_OK { throw WaifuError(code: code, message: lastError()) } + if code != WAIFU_MODELMANAGER_OK { throw WaifuError(code: code, message: modelManagerLastError()) } } /// Reads published models' manifests out of `path` where the download directory has none: the @@ -522,12 +668,12 @@ public enum ModelManager { public static func setBundledManifests(_ path: String?) throws { let code = path.map { $0.withCString { waifu_modelmanager_set_bundled_manifests($0) } } ?? waifu_modelmanager_set_bundled_manifests(nil) - if code != WAIFU_OK { throw WaifuError(code: code, message: lastError()) } + if code != WAIFU_MODELMANAGER_OK { throw WaifuError(code: code, message: modelManagerLastError()) } } public static func remove(_ name: String) throws { let code = name.withCString { waifu_modelmanager_remove($0) } - if code != WAIFU_OK { throw WaifuError(code: code, message: lastError()) } + if code != WAIFU_MODELMANAGER_OK { throw WaifuError(code: code, message: modelManagerLastError()) } } /// Fetches a published model without loading it. Throws CancellationError where it was @@ -540,9 +686,10 @@ public enum ModelManager { let fetched: Void? = try await call( progress: progress, cancel: { _ in waifu_modelmanager_fetch_cancel(handle.fetch) }, + refusal: modelManagerLastError, submit: { userData in handle.fetch = model.withCString { - waifu_modelmanager_fetch_async($0, userData, onProgress, onComplete) + waifu_modelmanager_fetch_async($0, userData, onFetchProgress, onFetchComplete) } return handle.fetch == nil ? 0 : 1 }) @@ -561,8 +708,9 @@ public enum ModelManager { let fetched: Void? = try await call( progress: progress, cancel: { _ in waifu_modelmanager_fetch_cancel(handle.fetch) }, + refusal: modelManagerLastError, submit: { userData in - handle.fetch = waifu_modelmanager_fetch_manifests_async(userData, onProgress, onComplete) + handle.fetch = waifu_modelmanager_fetch_manifests_async(userData, onFetchProgress, onFetchComplete) return handle.fetch == nil ? 0 : 1 }) waifu_modelmanager_fetch_free(handle.fetch) @@ -610,9 +758,20 @@ public enum Log { event.pointee.message.map { String(cString: $0) } ?? "") } - /// Every line the library would print goes to `handler` from now on, on whichever thread + /// The same for libwaifu_modelmanager, whose event is its own type. + private static let managerCallback: WaifuModelManagerLogCallback = { userData, event in + guard let userData, let event else { return } + let box = Unmanaged.fromOpaque(userData).takeUnretainedValue() + let level = LogLevel(rawValue: Int(event.pointee.level.rawValue)) ?? .info + box.handler( + level, + event.pointee.source.map { String(cString: $0) } ?? "", + event.pointee.message.map { String(cString: $0) } ?? "") + } + + /// Every line either library would print goes to `handler` from now on, on whichever thread /// wrote it; nil puts them back on the console. Set it before anything else is asked of the - /// library, to hear what hardware it found. + /// libraries, to hear what hardware libwaifu found. public static func setHandler(_ handler: LogHandler?) { lock.lock() defer { lock.unlock() } @@ -620,9 +779,11 @@ public enum Log { if let handler { let box = Unmanaged.passRetained(Box(handler)) waifu_set_log_callback(box.toOpaque(), callback) + waifu_modelmanager_set_log_callback(box.toOpaque(), managerCallback) current = box } else { waifu_set_log_callback(nil, nil) + waifu_modelmanager_set_log_callback(nil, nil) current = nil } // After the library has stopped calling it. @@ -631,5 +792,6 @@ public enum Log { public static func setLevel(_ level: LogLevel) { waifu_set_log_level(WaifuLogLevel(rawValue: UInt32(level.rawValue))) + waifu_modelmanager_set_log_level(WaifuModelManagerLogLevel(rawValue: UInt32(level.rawValue))) } } diff --git a/macos/Tests/WaifuKitTests/WaifuKitTests.swift b/waifu-app/swiftui/Tests/WaifuKitTests/WaifuKitTests.swift similarity index 88% rename from macos/Tests/WaifuKitTests/WaifuKitTests.swift rename to waifu-app/swiftui/Tests/WaifuKitTests/WaifuKitTests.swift index 9af869f..8f1b504 100644 --- a/macos/Tests/WaifuKitTests/WaifuKitTests.swift +++ b/waifu-app/swiftui/Tests/WaifuKitTests/WaifuKitTests.swift @@ -26,6 +26,7 @@ import XCTest final class WaifuKitTests: XCTestCase { func testTheLibraryIsTheOneTheHeaderDescribes() { XCTAssertEqual(ABI.library, ABI.header) + XCTAssertEqual(ABI.modelManagerLibrary, ABI.modelManagerHeader) } func testAVoiceIsLoadedAndSpeaks() async throws { @@ -115,8 +116,17 @@ final class WaifuKitTests: XCTestCase { // Every model says whether its manifest is here, which is what the app locks on. XCTAssertTrue(models.allSatisfy { $0["manifest_here"] is Bool }) - let described = try JSONSerialization.jsonObject(with: ModelManager.describeJSON("sdxl:base")) as? [String: Any] + // A manifest on the disk, whose weights are not: what it is, it says all the same. + let manifest = FileManager.default.temporaryDirectory + .appendingPathComponent("waifukit-\(ProcessInfo.processInfo.processIdentifier).yaml") + try "weights:\n - w.safetensors\nconfig:\n model:\n type: sdxl\n" + .write(to: manifest, atomically: true, encoding: .utf8) + defer { try? FileManager.default.removeItem(at: manifest) } + let described = try JSONSerialization.jsonObject(with: ModelManager.describeJSON(manifest.path)) as? [String: Any] XCTAssertEqual(described?["kind"] as? String, "image") + // What the boxes start at is the model manager's half of it, laid in. + XCTAssertEqual(described?["steps"] as? Int, 30) + XCTAssertEqual(described?["takes_guidance"] as? Bool, true) XCTAssertThrowsError(try ModelManager.describeJSON("no-such-model:v9")) XCTAssertNotNil(MachineInfo.json()) diff --git a/macos/build-app.sh b/waifu-app/swiftui/build-app.sh similarity index 76% rename from macos/build-app.sh rename to waifu-app/swiftui/build-app.sh index c78e763..e0733ee 100755 --- a/macos/build-app.sh +++ b/waifu-app/swiftui/build-app.sh @@ -1,11 +1,13 @@ #!/bin/sh -# Builds libwaifu-gui.app: the SwiftUI app, with libwaifu.dylib in its Frameworks, where the app's -# rpath looks for it. The executable is called libwaifu-gui, as the app is. +# Builds libwaifu-gui.app: the SwiftUI app, with libwaifu.dylib and libwaifu_modelmanager.dylib in +# its Frameworks, where the app's rpath looks for them. The executable is called libwaifu-gui, as +# the app is. # -# macos/build-app.sh [path/to/libwaifu.dylib] +# waifu-app/swiftui/build-app.sh [path/to/libwaifu.dylib] # # The library defaults to build/libwaifu.dylib, which the CMake build at the top of the repository -# makes (its waifu-ffi target). The app is written to macos/build/libwaifu-gui.app. +# makes (its waifu-engine target); libwaifu_modelmanager.dylib is taken from beside it (the +# waifu-modelmanager target). The app is written to waifu-app/swiftui/build/libwaifu-gui.app. # # WAIFU_SIGN_IDENTITY="Developer ID Application: Name (TEAMID)" signs it with that identity and the # hardened runtime, which notarization asks for; without it the app is signed for this machine only. @@ -17,11 +19,16 @@ set -eu here=$(cd "$(dirname "$0")" && pwd) -library=${1:-"$here/../build/libwaifu.dylib"} +library=${1:-"$here/../../build/libwaifu.dylib"} if [ ! -f "$library" ]; then echo "no libwaifu.dylib at $library: build it with cmake first, or name it" >&2 exit 1 fi +manager="$(dirname "$library")/libwaifu_modelmanager.dylib" +if [ ! -f "$manager" ]; then + echo "no libwaifu_modelmanager.dylib beside $library: build it with cmake first" >&2 + exit 1 +fi swift build --package-path "$here" -c release binary=$(swift build --package-path "$here" -c release --show-bin-path)/Waifu @@ -30,14 +37,15 @@ binary=$(swift build --package-path "$here" -c release --show-bin-path)/Waifu # app knows what each model suggests the moment it is chosen, with nothing to download first. manifests="$here/build/manifests" rm -rf "$manifests" -LIBWAIFU_LIB_DIR="$(cd "$here/.." && pwd)/build" cargo run --quiet --release \ - --manifest-path "$here/../waifu/Cargo.toml" --features hub --example fetch_manifests -- "$manifests" +LIBWAIFU_LIB_DIR="$(cd "$here/../.." && pwd)/build" cargo run --quiet --release \ + --manifest-path "$here/../../waifu-tools/Cargo.toml" --bin fetch_manifests -- "$manifests" app="$here/build/libwaifu-gui.app" rm -rf "$app" mkdir -p "$app/Contents/MacOS" "$app/Contents/Frameworks" "$app/Contents/Resources" cp "$binary" "$app/Contents/MacOS/libwaifu-gui" cp "$library" "$app/Contents/Frameworks/libwaifu.dylib" +cp "$manager" "$app/Contents/Frameworks/libwaifu_modelmanager.dylib" cp -R "$manifests" "$app/Contents/Resources/manifests" # The icon, in the format macOS 26 draws as its own rather than on a grey plate: an Icon Composer # .icon, compiled into Assets.car, with an AppIcon.icns beside it for older systems. @@ -77,7 +85,7 @@ EOF # The rpath to build/ is for running unbundled; inside the app the library is in Frameworks, and a # copy given away has no business looking in a folder on the machine that built it. -install_name_tool -delete_rpath "$(cd "$here/.." && pwd)/build" "$app/Contents/MacOS/libwaifu-gui" +install_name_tool -delete_rpath "$(cd "$here/../.." && pwd)/build" "$app/Contents/MacOS/libwaifu-gui" # Signed from the inside out, the library before the app that seals it. With no identity, for this # machine only, which is what lets it run here without a developer identity. @@ -87,6 +95,7 @@ else set -- --force --sign - fi codesign "$@" "$app/Contents/Frameworks/libwaifu.dylib" +codesign "$@" "$app/Contents/Frameworks/libwaifu_modelmanager.dylib" codesign "$@" "$app" codesign --verify --strict "$app" echo "$app" diff --git a/macos/dmg_layout.py b/waifu-app/swiftui/dmg_layout.py similarity index 100% rename from macos/dmg_layout.py rename to waifu-app/swiftui/dmg_layout.py diff --git a/macos/package-dmg.sh b/waifu-app/swiftui/package-dmg.sh similarity index 81% rename from macos/package-dmg.sh rename to waifu-app/swiftui/package-dmg.sh index 17ee436..07a86b4 100755 --- a/macos/package-dmg.sh +++ b/waifu-app/swiftui/package-dmg.sh @@ -2,13 +2,13 @@ # Packs libwaifu-gui.app as libwaifu-gui.dmg: one window, in icon view at 128 pixels, with the app # beside an Applications link to drag it onto. # -# macos/package-dmg.sh [path/to/libwaifu-gui.app] +# waifu-app/swiftui/package-dmg.sh [path/to/libwaifu-gui.app] # -# The app defaults to macos/build/libwaifu-gui.app, which build-app.sh makes. The image is written -# beside it, to macos/build/libwaifu-gui.dmg. +# The app defaults to waifu-app/swiftui/build/libwaifu-gui.app, which build-app.sh makes. The image is written +# beside it, to waifu-app/swiftui/build/libwaifu-gui.dmg. # # The window's layout is a .DS_Store that dmg_layout.py writes with the ds_store package, installed -# the first time into a venv under macos/build, so nothing here needs Finder or a desktop session. +# the first time into a venv under waifu-app/swiftui/build, so nothing here needs Finder or a desktop session. # # To hand it out: WAIFU_SIGN_IDENTITY, the Developer ID build-app.sh signed the app with, signs the # image too, and WAIFU_NOTARY_PROFILE, a profile `xcrun notarytool store-credentials` saved to the @@ -19,7 +19,7 @@ set -eu here=$(cd "$(dirname "$0")" && pwd) app=${1:-"$here/build/libwaifu-gui.app"} if [ ! -d "$app" ]; then - echo "no libwaifu-gui.app at $app: run macos/build-app.sh first, or name it" >&2 + echo "no libwaifu-gui.app at $app: run waifu-app/swiftui/build-app.sh first, or name it" >&2 exit 1 fi codesign --verify --deep "$app" diff --git a/windows/.gitignore b/waifu-app/wpf/.gitignore similarity index 100% rename from windows/.gitignore rename to waifu-app/wpf/.gitignore diff --git a/windows/Libwaifu.Gui.sln b/waifu-app/wpf/Libwaifu.Gui.sln similarity index 100% rename from windows/Libwaifu.Gui.sln rename to waifu-app/wpf/Libwaifu.Gui.sln diff --git a/windows/Libwaifu.Gui/App.config b/waifu-app/wpf/Libwaifu.Gui/App.config similarity index 100% rename from windows/Libwaifu.Gui/App.config rename to waifu-app/wpf/Libwaifu.Gui/App.config diff --git a/windows/Libwaifu.Gui/App.xaml b/waifu-app/wpf/Libwaifu.Gui/App.xaml similarity index 100% rename from windows/Libwaifu.Gui/App.xaml rename to waifu-app/wpf/Libwaifu.Gui/App.xaml diff --git a/windows/Libwaifu.Gui/App.xaml.cs b/waifu-app/wpf/Libwaifu.Gui/App.xaml.cs similarity index 100% rename from windows/Libwaifu.Gui/App.xaml.cs rename to waifu-app/wpf/Libwaifu.Gui/App.xaml.cs diff --git a/windows/Libwaifu.Gui/Controls/DropBox.cs b/waifu-app/wpf/Libwaifu.Gui/Controls/DropBox.cs similarity index 100% rename from windows/Libwaifu.Gui/Controls/DropBox.cs rename to waifu-app/wpf/Libwaifu.Gui/Controls/DropBox.cs diff --git a/windows/Libwaifu.Gui/Controls/FormGroup.cs b/waifu-app/wpf/Libwaifu.Gui/Controls/FormGroup.cs similarity index 100% rename from windows/Libwaifu.Gui/Controls/FormGroup.cs rename to waifu-app/wpf/Libwaifu.Gui/Controls/FormGroup.cs diff --git a/windows/Libwaifu.Gui/Controls/SettingRow.cs b/waifu-app/wpf/Libwaifu.Gui/Controls/SettingRow.cs similarity index 100% rename from windows/Libwaifu.Gui/Controls/SettingRow.cs rename to waifu-app/wpf/Libwaifu.Gui/Controls/SettingRow.cs diff --git a/windows/Libwaifu.Gui/Converters/Converters.cs b/waifu-app/wpf/Libwaifu.Gui/Converters/Converters.cs similarity index 100% rename from windows/Libwaifu.Gui/Converters/Converters.cs rename to waifu-app/wpf/Libwaifu.Gui/Converters/Converters.cs diff --git a/windows/Libwaifu.Gui/Engine/Catalog.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/Catalog.cs similarity index 99% rename from windows/Libwaifu.Gui/Engine/Catalog.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/Catalog.cs index f33a7c2..72f9c0c 100644 --- a/windows/Libwaifu.Gui/Engine/Catalog.cs +++ b/waifu-app/wpf/Libwaifu.Gui/Engine/Catalog.cs @@ -55,7 +55,7 @@ public class ModelEntry /// What a model says about itself: the numbers to start from and the things it cannot do. /// /// - /// This is what the real screen gets back from waifu_modelmanager_describe_json, and + /// This is what the real screen gets back from waifu_describe_json, and /// choosing a model asks for it and nothing else -- reading the description is not loading /// the model, and no job runs until the Generate button is pressed. /// diff --git a/windows/Libwaifu.Gui/Engine/Core.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/Core.cs similarity index 100% rename from windows/Libwaifu.Gui/Engine/Core.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/Core.cs diff --git a/windows/Libwaifu.Gui/Engine/Json.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/Json.cs similarity index 100% rename from windows/Libwaifu.Gui/Engine/Json.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/Json.cs diff --git a/windows/Libwaifu.Gui/Engine/Library.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/Library.cs similarity index 76% rename from windows/Libwaifu.Gui/Engine/Library.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/Library.cs index cdcf987..5798f7f 100644 --- a/windows/Libwaifu.Gui/Engine/Library.cs +++ b/waifu-app/wpf/Libwaifu.Gui/Engine/Library.cs @@ -38,6 +38,7 @@ public static class Library // The library keeps this pointer. A delegate that is collected is a crash in somebody // else's stack frame, so it is held here for as long as the process lives. private static Native.LogCallback logSink; + private static Native.ManagerLogCallback managerLogSink; private static bool looked; private static bool here; @@ -88,6 +89,26 @@ private static bool Open() return false; } + // The model manager is a library of its own, and the screen needs both: it is + // looked for beside the engine's, where the build writes the two of them. + var manager = System.IO.Path.Combine( + System.IO.Path.GetDirectoryName(found), Native.ModelManagerLibrary); + if (!File.Exists(manager)) + { + Missing = Native.ModelManagerLibrary + " is not beside " + found + "."; + Log.Write("error", "library", Missing); + return false; + } + + if (Native.LoadLibrary(manager) == IntPtr.Zero) + { + Missing = "could not open " + manager + ": " + + new System.ComponentModel.Win32Exception( + System.Runtime.InteropServices.Marshal.GetLastWin32Error()).Message; + Log.Write("error", "library", Missing); + return false; + } + Path = found; var abi = Native.waifu_abi_version(); @@ -99,6 +120,15 @@ private static bool Open() return false; } + var managerAbi = Native.waifu_modelmanager_abi_version(); + if (managerAbi != Native.ModelManagerAbiVersion) + { + Missing = manager + " is ABI " + managerAbi + " and this program was built for " + + Native.ModelManagerAbiVersion + "."; + Log.Write("error", "library", Missing); + return false; + } + Listen(); Bundled(); Log.Write("info", "library", "opened " + found + ", ABI " + abi); @@ -141,17 +171,17 @@ private static System.Collections.Generic.IEnumerable Folders() yield return beside; yield return Environment.GetEnvironmentVariable("LIBWAIFU_LIB_DIR"); - // ...\windows\Libwaifu.Gui\bin\Debug -> the repository's build directory. + // ...\waifu-app\wpf\Libwaifu.Gui\bin\Debug -> the repository's build directory. if (beside != null) { var repository = System.IO.Path.GetFullPath( - System.IO.Path.Combine(beside, "..", "..", "..", "..")); + System.IO.Path.Combine(beside, "..", "..", "..", "..", "..")); yield return System.IO.Path.Combine(repository, "build"); } } /// - /// Every line the library would print, from whichever thread wrote it. Set before + /// Every line either library would print, from whichever thread wrote it. Set before /// anything else: flint writes what hardware it found the first time a device is asked /// about, and that line is worth having. /// @@ -160,6 +190,9 @@ private static void Listen() logSink = OnLine; Native.waifu_set_log_callback(IntPtr.Zero, logSink); Native.waifu_set_log_level(Native.LogLevel.Info); + managerLogSink = OnManagerLine; + Native.waifu_modelmanager_set_log_callback(IntPtr.Zero, managerLogSink); + Native.waifu_modelmanager_set_log_level(Native.LogLevel.Info); } private static void OnLine(IntPtr user, ref Native.LogEvent e) @@ -177,6 +210,19 @@ private static void OnLine(IntPtr user, ref Native.LogEvent e) } } + /// The same for waifu_modelmanager.dll, whose event is its own type. + private static void OnManagerLine(IntPtr user, ref Native.ManagerLogEvent e) + { + try + { + Log.Write(Level(e.Level), Utf8.Read(e.Source) ?? "models", Utf8.Read(e.Message)); + } + catch + { + // As above: nowhere for an exception to go. + } + } + private static string Level(Native.LogLevel level) { switch (level) @@ -217,6 +263,19 @@ public static string LastError() } } + /// The same for a call into the model manager, which keeps its own. + public static string ModelManagerLastError() + { + try + { + return Utf8.Read(Native.waifu_modelmanager_last_error()) ?? "no reason given"; + } + catch (Exception error) + { + return error.Message; + } + } + internal static Native.Device DeviceNamed(string name) { switch ((name ?? "auto").Trim().ToLowerInvariant()) diff --git a/windows/Libwaifu.Gui/Engine/Machine.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/Machine.cs similarity index 100% rename from windows/Libwaifu.Gui/Engine/Machine.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/Machine.cs diff --git a/windows/Libwaifu.Gui/Engine/MockEngine.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/MockEngine.cs similarity index 100% rename from windows/Libwaifu.Gui/Engine/MockEngine.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/MockEngine.cs diff --git a/windows/Libwaifu.Gui/Engine/Models.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/Models.cs similarity index 100% rename from windows/Libwaifu.Gui/Engine/Models.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/Models.cs diff --git a/windows/Libwaifu.Gui/Engine/Native.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/Native.cs similarity index 70% rename from windows/Libwaifu.Gui/Engine/Native.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/Native.cs index f27c4fa..9f14a19 100644 --- a/windows/Libwaifu.Gui/Engine/Native.cs +++ b/waifu-app/wpf/Libwaifu.Gui/Engine/Native.cs @@ -46,8 +46,17 @@ internal static class Native /// public const string Library = "waifu.dll"; + /// + /// The model manager: the catalogue, fetching, where models are kept. A library of its + /// own, which links no engine, beside the engine's. + /// + public const string ModelManagerLibrary = "waifu_modelmanager.dll"; + /// What a caller compiled against this header expects, from waifu.h. - public const uint AbiVersion = 1; + public const uint AbiVersion = 2; + + /// The same, from waifu_modelmanager.h. + public const uint ModelManagerAbiVersion = 1; public enum LogLevel { @@ -115,6 +124,63 @@ public struct CompletionEvent public StatusInfo Status; } + // -- waifu_modelmanager.dll's own types, which share nothing with the above ---------- + + public enum ManagerStatus + { + Ok = 0, + Cancelled = 1, + InvalidArgument = 2, + UnknownModel = 3, + Fetch = 4, + Internal = 5, + } + + [StructLayout(LayoutKind.Sequential)] + public struct ManagerLogEvent + { + public LogLevel Level; + public IntPtr Source; + public IntPtr Message; + } + + /// How far a download is: the file, its bytes, and which hub it comes from. + [StructLayout(LayoutKind.Sequential)] + public struct ManagerProgressEvent + { + public double Fraction; + public ulong Done; + public ulong Total; + public uint Part; + public uint Parts; + public IntPtr File; + public IntPtr Words; + public double Seconds; + public IntPtr Hub; + } + + [StructLayout(LayoutKind.Sequential)] + public struct ManagerStatusInfo + { + public ManagerStatus Code; + public IntPtr Message; + } + + [StructLayout(LayoutKind.Sequential)] + public struct ManagerCompletionEvent + { + public ManagerStatusInfo Status; + } + + [UnmanagedFunctionPointer(CallingConvention.Cdecl)] + public delegate void ManagerLogCallback(IntPtr user, ref ManagerLogEvent e); + + [UnmanagedFunctionPointer(CallingConvention.Cdecl)] + public delegate void ManagerProgressCallback(IntPtr user, ref ManagerProgressEvent e); + + [UnmanagedFunctionPointer(CallingConvention.Cdecl)] + public delegate void ManagerCompleteCallback(IntPtr user, ref ManagerCompletionEvent e); + [StructLayout(LayoutKind.Sequential)] public struct Image { @@ -230,36 +296,58 @@ public struct VoiceConversionRequest [DllImport(Library, CallingConvention = CallingConvention.Cdecl)] public static extern IntPtr waifu_machine_json(); + /// What the manifest at a path says, or "tones": see waifu_modelmanager_manifest_path. [DllImport(Library, CallingConvention = CallingConvention.Cdecl)] - public static extern IntPtr waifu_modelmanager_catalog_json(); + public static extern IntPtr waifu_describe_json(byte[] model); - [DllImport(Library, CallingConvention = CallingConvention.Cdecl)] - public static extern IntPtr waifu_modelmanager_describe_json(byte[] model); + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] + public static extern uint waifu_modelmanager_abi_version(); - [DllImport(Library, CallingConvention = CallingConvention.Cdecl)] + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] + public static extern IntPtr waifu_modelmanager_last_error(); + + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] + public static extern void waifu_modelmanager_string_free(IntPtr text); + + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] + public static extern void waifu_modelmanager_set_log_callback(IntPtr user, ManagerLogCallback callback); + + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] + public static extern void waifu_modelmanager_set_log_level(LogLevel level); + + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] + public static extern IntPtr waifu_modelmanager_catalog_json(); + + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] public static extern IntPtr waifu_modelmanager_directory(); - [DllImport(Library, CallingConvention = CallingConvention.Cdecl)] - public static extern Status waifu_modelmanager_set_directory(byte[] path); + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] + public static extern ManagerStatus waifu_modelmanager_set_directory(byte[] path); - [DllImport(Library, CallingConvention = CallingConvention.Cdecl)] - public static extern Status waifu_modelmanager_set_bundled_manifests(byte[] directory); + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] + public static extern ManagerStatus waifu_modelmanager_set_bundled_manifests(byte[] directory); - [DllImport(Library, CallingConvention = CallingConvention.Cdecl)] - public static extern Status waifu_modelmanager_remove(byte[] name); + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] + public static extern ManagerStatus waifu_modelmanager_remove(byte[] name); - [DllImport(Library, CallingConvention = CallingConvention.Cdecl)] + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] + public static extern IntPtr waifu_modelmanager_manifest_path(byte[] name); + + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] + public static extern IntPtr waifu_modelmanager_defaults_json(byte[] model); + + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] public static extern IntPtr waifu_modelmanager_fetch_async( - byte[] model, IntPtr user, ProgressCallback onProgress, CompleteCallback onComplete); + byte[] model, IntPtr user, ManagerProgressCallback onProgress, ManagerCompleteCallback onComplete); - [DllImport(Library, CallingConvention = CallingConvention.Cdecl)] + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] public static extern IntPtr waifu_modelmanager_fetch_manifests_async( - IntPtr user, ProgressCallback onProgress, CompleteCallback onComplete); + IntPtr user, ManagerProgressCallback onProgress, ManagerCompleteCallback onComplete); - [DllImport(Library, CallingConvention = CallingConvention.Cdecl)] + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] public static extern void waifu_modelmanager_fetch_cancel(IntPtr fetch); - [DllImport(Library, CallingConvention = CallingConvention.Cdecl)] + [DllImport(ModelManagerLibrary, CallingConvention = CallingConvention.Cdecl)] public static extern void waifu_modelmanager_fetch_free(IntPtr fetch); [DllImport(Library, CallingConvention = CallingConvention.Cdecl)] @@ -379,5 +467,23 @@ public static string Take(IntPtr pointer) Native.waifu_string_free(pointer); } } + + /// The same for a string the model manager returned, which it frees itself. + public static string TakeFromModelManager(IntPtr pointer) + { + if (pointer == IntPtr.Zero) + { + return null; + } + + try + { + return Read(pointer); + } + finally + { + Native.waifu_modelmanager_string_free(pointer); + } + } } } diff --git a/windows/Libwaifu.Gui/Engine/NativeEngine.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/NativeEngine.cs similarity index 90% rename from windows/Libwaifu.Gui/Engine/NativeEngine.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/NativeEngine.cs index 147f2b5..8926e43 100644 --- a/windows/Libwaifu.Gui/Engine/NativeEngine.cs +++ b/waifu-app/wpf/Libwaifu.Gui/Engine/NativeEngine.cs @@ -55,6 +55,7 @@ public sealed class NativeEngine : IEngine private IntPtr engine; private string device; private ulong job; + private Fetch fetching; private JobRequest current; private long seed; private double started; @@ -95,26 +96,83 @@ public void Begin(JobRequest request) return; } - // Loading fetches what is not here and drops whatever was loaded before it, so there - // is nothing to unload first. - Say(Stage.Reading, null, "reading " + request.Model); - job = Native.waifu_engine_load_async( - engine, Utf8.Bytes(request.Model), IntPtr.Zero, onProgress, onLoaded); - - if (job == 0) + // waifu.dll reads a model by the path of its manifest. A published one not all here + // is fetched by the model manager first, and its manifest found there. + if (NativeModels.NeedsFetching(request.Model)) { - End(new JobOutcome { Ending = Ending.Failed, Error = Library.LastError() }); + FetchThenLoad(request); + return; } + + Load(request); } public void Cancel() { + if (fetching != null) + { + fetching.Cancel(); + } + if (engine != IntPtr.Zero && job != 0) { Native.waifu_engine_cancel(engine, job); } } + /// Fetches the model a job needs, and loads it once it is here. + private void FetchThenLoad(JobRequest request) + { + fetching = Fetch.Start(request.Model); + if (fetching == null) + { + End(new JobOutcome { Ending = Ending.Failed, Error = Library.ModelManagerLastError() }); + return; + } + + fetching.Progress += (fraction, line) => + Say(Stage.Fetching, fraction >= 0 ? (double?)fraction : null, line); + fetching.Finished += error => + { + fetching = null; + if (error == null) + { + Load(request); + } + else if (error.Length == 0) + { + End(new JobOutcome { Ending = Ending.Cancelled }); + } + else + { + End(new JobOutcome { Ending = Ending.Failed, Error = error }); + } + }; + } + + /// + /// Reads the model onto the device, dropping whatever was loaded before it, so there is + /// nothing to unload first. + /// + private void Load(JobRequest request) + { + var manifest = NativeModels.ManifestPath(request.Model); + if (manifest == null) + { + End(new JobOutcome { Ending = Ending.Failed, Error = Library.ModelManagerLastError() }); + return; + } + + Say(Stage.Reading, null, "reading " + request.Model); + job = Native.waifu_engine_load_async( + engine, Utf8.Bytes(manifest), IntPtr.Zero, onProgress, onLoaded); + + if (job == 0) + { + End(new JobOutcome { Ending = Ending.Failed, Error = Library.LastError() }); + } + } + public void Dispose() { Drop(); diff --git a/windows/Libwaifu.Gui/Engine/NativeModels.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/NativeModels.cs similarity index 73% rename from windows/Libwaifu.Gui/Engine/NativeModels.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/NativeModels.cs index 5d22e29..614d098 100644 --- a/windows/Libwaifu.Gui/Engine/NativeModels.cs +++ b/waifu-app/wpf/Libwaifu.Gui/Engine/NativeModels.cs @@ -19,6 +19,7 @@ using System; using System.Collections.Generic; +using System.IO; using System.Linq; using System.Windows.Threading; @@ -34,10 +35,10 @@ public static class NativeModels public static List Catalog() { var found = new List(); - var text = Utf8.Take(Native.waifu_modelmanager_catalog_json()); + var text = Utf8.TakeFromModelManager(Native.waifu_modelmanager_catalog_json()); if (text == null) { - Log.Write("warn", "models", "no catalog: " + Library.LastError()); + Log.Write("warn", "models", "no catalog: " + Library.ModelManagerLastError()); return found; } @@ -63,24 +64,74 @@ public static List Catalog() .ToList(); } + /// The stand-in voice, which is in waifu.dll itself and has no manifest. + public const string Tones = "tones"; + + /// + /// The manifest of a published model here -- fetched, or bundled -- or a manifest's own + /// path, as it is: what waifu.dll is handed. Null where there is none here, with the + /// reason in . + /// + public static string ManifestPath(string model) + { + if (model == Tones) + { + return model; + } + + return Utf8.TakeFromModelManager(Native.waifu_modelmanager_manifest_path(Utf8.Bytes(model))); + } + + /// What the catalogue calls a published model, or null for anything else. + public static string PublishedFullName(string model) + { + var listed = Catalog().FirstOrDefault(entry => entry.Name == model); + return listed == null ? null : listed.FullName; + } + + /// + /// Whether a load of has a download to do first: a published + /// model not all here. The stand-in and a manifest on the disk have nothing to fetch. + /// + public static bool NeedsFetching(string model) + { + if (model == Tones || File.Exists(model)) + { + return false; + } + + var listed = Catalog().FirstOrDefault(entry => entry.Name == model); + return listed != null && !listed.Cached; + } + /// - /// What a model is, without reading it: a published name or the path of a manifest. Null - /// where the library cannot say, with the reason in . + /// What a model is, without reading it: what waifu.dll says of the manifest the model + /// manager finds for it, whose weights need not be here. Null where there is no manifest + /// here, or the library cannot say. /// public static ModelDescription Describe(string model) { - var text = Utf8.Take(Native.waifu_modelmanager_describe_json(Utf8.Bytes(model))); + var manifest = ManifestPath(model); + if (manifest == null) + { + return null; + } + + var text = Utf8.Take(Native.waifu_describe_json(Utf8.Bytes(manifest))); if (text == null) { return null; } var node = Json.Parse(text); + WithDefaults(node, manifest); var kind = Json.Text(node, "kind"); var description = new ModelDescription { - Name = Json.Text(node, "name") ?? model, - FullName = Json.Text(node, "full_name") ?? Json.Text(node, "name") ?? model, + Name = model, + // waifu.dll calls a model by its manifest's file name; one asked for by its published + // name is called what the catalogue calls it. + FullName = PublishedFullName(model) ?? Json.Text(node, "full_name") ?? model, Category = CategoryNamed(kind), OnDisk = Json.Flag(node, "on_disk"), InMemory = Json.Flag(node, "in_memory"), @@ -102,6 +153,40 @@ public static ModelDescription Describe(string model) return description; } + /// + /// Lays what waifu_modelmanager.dll says a model's boxes start at into what waifu.dll says + /// it can be asked for: the two halves of one description. A "conversion" is in both, and + /// its defaults go inside it. + /// + private static void WithDefaults(object node, string manifest) + { + var text = Utf8.TakeFromModelManager( + Native.waifu_modelmanager_defaults_json(Utf8.Bytes(manifest))); + if (text == null) + { + Log.Write("warn", "models", manifest + ": " + Library.ModelManagerLastError()); + return; + } + + var whole = Json.Map(node); + foreach (var field in Json.Map(Json.Parse(text))) + { + var inner = whole.ContainsKey(field.Key) ? whole[field.Key] as Dictionary : null; + var more = field.Value as Dictionary; + if (inner != null && more != null) + { + foreach (var each in more) + { + inner[each.Key] = each.Value; + } + } + else + { + whole[field.Key] = field.Value; + } + } + } + private static void ReadPicture(object node, ModelDescription description) { description.Steps = Json.Count(node, "steps", 20); @@ -177,22 +262,22 @@ public static Category CategoryNamed(string kind) public static string Directory() { - return Utf8.Take(Native.waifu_modelmanager_directory()); + return Utf8.TakeFromModelManager(Native.waifu_modelmanager_directory()); } public static bool SetDirectory(string path) { - return Native.waifu_modelmanager_set_directory(Utf8.Bytes(path)) == Native.Status.Ok; + return Native.waifu_modelmanager_set_directory(Utf8.Bytes(path)) == Native.ManagerStatus.Ok; } public static bool Remove(string name) { - if (Native.waifu_modelmanager_remove(Utf8.Bytes(name)) == Native.Status.Ok) + if (Native.waifu_modelmanager_remove(Utf8.Bytes(name)) == Native.ManagerStatus.Ok) { return true; } - Log.Write("error", "models", name + ": " + Library.LastError()); + Log.Write("error", "models", name + ": " + Library.ModelManagerLastError()); return false; } @@ -253,8 +338,8 @@ public static MachineSnapshot Machine() public sealed class Fetch : IDisposable { private readonly Dispatcher dispatcher; - private readonly Native.ProgressCallback onProgress; - private readonly Native.CompleteCallback onComplete; + private readonly Native.ManagerProgressCallback onProgress; + private readonly Native.ManagerCompleteCallback onComplete; private IntPtr handle; private bool ended; @@ -285,7 +370,7 @@ public static Fetch Start(string model) if (fetch.handle == IntPtr.Zero) { - Log.Write("error", "models", model + ": " + Library.LastError()); + Log.Write("error", "models", model + ": " + Library.ModelManagerLastError()); return null; } @@ -314,7 +399,7 @@ public void Dispose() } } - private void Advanced(IntPtr user, ref Native.ProgressEvent e) + private void Advanced(IntPtr user, ref Native.ManagerProgressEvent e) { var fraction = e.Fraction; var words = Utf8.Read(e.Words); @@ -331,7 +416,7 @@ private void Advanced(IntPtr user, ref Native.ProgressEvent e) }); } - private void Ended(IntPtr user, ref Native.CompletionEvent e) + private void Ended(IntPtr user, ref Native.ManagerCompletionEvent e) { var code = e.Status.Code; var message = Utf8.Read(e.Status.Message); @@ -347,9 +432,9 @@ private void Ended(IntPtr user, ref Native.CompletionEvent e) var handler = Finished; if (handler != null) { - handler(code == Native.Status.Ok + handler(code == Native.ManagerStatus.Ok ? null - : (code == Native.Status.Cancelled ? string.Empty : message ?? "it failed")); + : (code == Native.ManagerStatus.Cancelled ? string.Empty : message ?? "it failed")); } Dispose(); diff --git a/windows/Libwaifu.Gui/Engine/Output.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/Output.cs similarity index 100% rename from windows/Libwaifu.Gui/Engine/Output.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/Output.cs diff --git a/windows/Libwaifu.Gui/Engine/Sound.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/Sound.cs similarity index 100% rename from windows/Libwaifu.Gui/Engine/Sound.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/Sound.cs diff --git a/windows/Libwaifu.Gui/Engine/Store.cs b/waifu-app/wpf/Libwaifu.Gui/Engine/Store.cs similarity index 100% rename from windows/Libwaifu.Gui/Engine/Store.cs rename to waifu-app/wpf/Libwaifu.Gui/Engine/Store.cs diff --git a/windows/Libwaifu.Gui/Libwaifu.Gui.csproj b/waifu-app/wpf/Libwaifu.Gui/Libwaifu.Gui.csproj similarity index 100% rename from windows/Libwaifu.Gui/Libwaifu.Gui.csproj rename to waifu-app/wpf/Libwaifu.Gui/Libwaifu.Gui.csproj diff --git a/windows/Libwaifu.Gui/Properties/AssemblyInfo.cs b/waifu-app/wpf/Libwaifu.Gui/Properties/AssemblyInfo.cs similarity index 100% rename from windows/Libwaifu.Gui/Properties/AssemblyInfo.cs rename to waifu-app/wpf/Libwaifu.Gui/Properties/AssemblyInfo.cs diff --git a/windows/Libwaifu.Gui/Theme/Controls.xaml b/waifu-app/wpf/Libwaifu.Gui/Theme/Controls.xaml similarity index 100% rename from windows/Libwaifu.Gui/Theme/Controls.xaml rename to waifu-app/wpf/Libwaifu.Gui/Theme/Controls.xaml diff --git a/windows/Libwaifu.Gui/Theme/Icons.xaml b/waifu-app/wpf/Libwaifu.Gui/Theme/Icons.xaml similarity index 100% rename from windows/Libwaifu.Gui/Theme/Icons.xaml rename to waifu-app/wpf/Libwaifu.Gui/Theme/Icons.xaml diff --git a/windows/Libwaifu.Gui/Theme/Tokens.xaml b/waifu-app/wpf/Libwaifu.Gui/Theme/Tokens.xaml similarity index 100% rename from windows/Libwaifu.Gui/Theme/Tokens.xaml rename to waifu-app/wpf/Libwaifu.Gui/Theme/Tokens.xaml diff --git a/windows/Libwaifu.Gui/ViewModels/Common.cs b/waifu-app/wpf/Libwaifu.Gui/ViewModels/Common.cs similarity index 100% rename from windows/Libwaifu.Gui/ViewModels/Common.cs rename to waifu-app/wpf/Libwaifu.Gui/ViewModels/Common.cs diff --git a/windows/Libwaifu.Gui/ViewModels/ImageTaskViewModel.cs b/waifu-app/wpf/Libwaifu.Gui/ViewModels/ImageTaskViewModel.cs similarity index 100% rename from windows/Libwaifu.Gui/ViewModels/ImageTaskViewModel.cs rename to waifu-app/wpf/Libwaifu.Gui/ViewModels/ImageTaskViewModel.cs diff --git a/windows/Libwaifu.Gui/ViewModels/ModelManagerViewModel.cs b/waifu-app/wpf/Libwaifu.Gui/ViewModels/ModelManagerViewModel.cs similarity index 100% rename from windows/Libwaifu.Gui/ViewModels/ModelManagerViewModel.cs rename to waifu-app/wpf/Libwaifu.Gui/ViewModels/ModelManagerViewModel.cs diff --git a/windows/Libwaifu.Gui/ViewModels/Shell.cs b/waifu-app/wpf/Libwaifu.Gui/ViewModels/Shell.cs similarity index 100% rename from windows/Libwaifu.Gui/ViewModels/Shell.cs rename to waifu-app/wpf/Libwaifu.Gui/ViewModels/Shell.cs diff --git a/windows/Libwaifu.Gui/ViewModels/SpeechTaskViewModel.cs b/waifu-app/wpf/Libwaifu.Gui/ViewModels/SpeechTaskViewModel.cs similarity index 100% rename from windows/Libwaifu.Gui/ViewModels/SpeechTaskViewModel.cs rename to waifu-app/wpf/Libwaifu.Gui/ViewModels/SpeechTaskViewModel.cs diff --git a/windows/Libwaifu.Gui/ViewModels/TaskViewModel.cs b/waifu-app/wpf/Libwaifu.Gui/ViewModels/TaskViewModel.cs similarity index 100% rename from windows/Libwaifu.Gui/ViewModels/TaskViewModel.cs rename to waifu-app/wpf/Libwaifu.Gui/ViewModels/TaskViewModel.cs diff --git a/windows/Libwaifu.Gui/Views/ConversionInspector.xaml b/waifu-app/wpf/Libwaifu.Gui/Views/ConversionInspector.xaml similarity index 100% rename from windows/Libwaifu.Gui/Views/ConversionInspector.xaml rename to waifu-app/wpf/Libwaifu.Gui/Views/ConversionInspector.xaml diff --git a/windows/Libwaifu.Gui/Views/ConversionInspector.xaml.cs b/waifu-app/wpf/Libwaifu.Gui/Views/ConversionInspector.xaml.cs similarity index 100% rename from windows/Libwaifu.Gui/Views/ConversionInspector.xaml.cs rename to waifu-app/wpf/Libwaifu.Gui/Views/ConversionInspector.xaml.cs diff --git a/windows/Libwaifu.Gui/Views/ConversionTaskView.xaml b/waifu-app/wpf/Libwaifu.Gui/Views/ConversionTaskView.xaml similarity index 100% rename from windows/Libwaifu.Gui/Views/ConversionTaskView.xaml rename to waifu-app/wpf/Libwaifu.Gui/Views/ConversionTaskView.xaml diff --git a/windows/Libwaifu.Gui/Views/ConversionTaskView.xaml.cs b/waifu-app/wpf/Libwaifu.Gui/Views/ConversionTaskView.xaml.cs similarity index 100% rename from windows/Libwaifu.Gui/Views/ConversionTaskView.xaml.cs rename to waifu-app/wpf/Libwaifu.Gui/Views/ConversionTaskView.xaml.cs diff --git a/windows/Libwaifu.Gui/Views/ImageInspector.xaml b/waifu-app/wpf/Libwaifu.Gui/Views/ImageInspector.xaml similarity index 100% rename from windows/Libwaifu.Gui/Views/ImageInspector.xaml rename to waifu-app/wpf/Libwaifu.Gui/Views/ImageInspector.xaml diff --git a/windows/Libwaifu.Gui/Views/ImageInspector.xaml.cs b/waifu-app/wpf/Libwaifu.Gui/Views/ImageInspector.xaml.cs similarity index 100% rename from windows/Libwaifu.Gui/Views/ImageInspector.xaml.cs rename to waifu-app/wpf/Libwaifu.Gui/Views/ImageInspector.xaml.cs diff --git a/windows/Libwaifu.Gui/Views/ImageTaskView.xaml b/waifu-app/wpf/Libwaifu.Gui/Views/ImageTaskView.xaml similarity index 100% rename from windows/Libwaifu.Gui/Views/ImageTaskView.xaml rename to waifu-app/wpf/Libwaifu.Gui/Views/ImageTaskView.xaml diff --git a/windows/Libwaifu.Gui/Views/ImageTaskView.xaml.cs b/waifu-app/wpf/Libwaifu.Gui/Views/ImageTaskView.xaml.cs similarity index 100% rename from windows/Libwaifu.Gui/Views/ImageTaskView.xaml.cs rename to waifu-app/wpf/Libwaifu.Gui/Views/ImageTaskView.xaml.cs diff --git a/windows/Libwaifu.Gui/Views/InspectorSelector.cs b/waifu-app/wpf/Libwaifu.Gui/Views/InspectorSelector.cs similarity index 100% rename from windows/Libwaifu.Gui/Views/InspectorSelector.cs rename to waifu-app/wpf/Libwaifu.Gui/Views/InspectorSelector.cs diff --git a/windows/Libwaifu.Gui/Views/LogWindow.xaml b/waifu-app/wpf/Libwaifu.Gui/Views/LogWindow.xaml similarity index 100% rename from windows/Libwaifu.Gui/Views/LogWindow.xaml rename to waifu-app/wpf/Libwaifu.Gui/Views/LogWindow.xaml diff --git a/windows/Libwaifu.Gui/Views/LogWindow.xaml.cs b/waifu-app/wpf/Libwaifu.Gui/Views/LogWindow.xaml.cs similarity index 100% rename from windows/Libwaifu.Gui/Views/LogWindow.xaml.cs rename to waifu-app/wpf/Libwaifu.Gui/Views/LogWindow.xaml.cs diff --git a/windows/Libwaifu.Gui/Views/MainWindow.xaml b/waifu-app/wpf/Libwaifu.Gui/Views/MainWindow.xaml similarity index 100% rename from windows/Libwaifu.Gui/Views/MainWindow.xaml rename to waifu-app/wpf/Libwaifu.Gui/Views/MainWindow.xaml diff --git a/windows/Libwaifu.Gui/Views/MainWindow.xaml.cs b/waifu-app/wpf/Libwaifu.Gui/Views/MainWindow.xaml.cs similarity index 100% rename from windows/Libwaifu.Gui/Views/MainWindow.xaml.cs rename to waifu-app/wpf/Libwaifu.Gui/Views/MainWindow.xaml.cs diff --git a/windows/Libwaifu.Gui/Views/ModelManagerView.xaml b/waifu-app/wpf/Libwaifu.Gui/Views/ModelManagerView.xaml similarity index 100% rename from windows/Libwaifu.Gui/Views/ModelManagerView.xaml rename to waifu-app/wpf/Libwaifu.Gui/Views/ModelManagerView.xaml diff --git a/windows/Libwaifu.Gui/Views/ModelManagerView.xaml.cs b/waifu-app/wpf/Libwaifu.Gui/Views/ModelManagerView.xaml.cs similarity index 100% rename from windows/Libwaifu.Gui/Views/ModelManagerView.xaml.cs rename to waifu-app/wpf/Libwaifu.Gui/Views/ModelManagerView.xaml.cs diff --git a/windows/Libwaifu.Gui/Views/PlayerView.xaml b/waifu-app/wpf/Libwaifu.Gui/Views/PlayerView.xaml similarity index 100% rename from windows/Libwaifu.Gui/Views/PlayerView.xaml rename to waifu-app/wpf/Libwaifu.Gui/Views/PlayerView.xaml diff --git a/windows/Libwaifu.Gui/Views/PlayerView.xaml.cs b/waifu-app/wpf/Libwaifu.Gui/Views/PlayerView.xaml.cs similarity index 100% rename from windows/Libwaifu.Gui/Views/PlayerView.xaml.cs rename to waifu-app/wpf/Libwaifu.Gui/Views/PlayerView.xaml.cs diff --git a/windows/Libwaifu.Gui/Views/SpeechInspector.xaml b/waifu-app/wpf/Libwaifu.Gui/Views/SpeechInspector.xaml similarity index 100% rename from windows/Libwaifu.Gui/Views/SpeechInspector.xaml rename to waifu-app/wpf/Libwaifu.Gui/Views/SpeechInspector.xaml diff --git a/windows/Libwaifu.Gui/Views/SpeechInspector.xaml.cs b/waifu-app/wpf/Libwaifu.Gui/Views/SpeechInspector.xaml.cs similarity index 100% rename from windows/Libwaifu.Gui/Views/SpeechInspector.xaml.cs rename to waifu-app/wpf/Libwaifu.Gui/Views/SpeechInspector.xaml.cs diff --git a/windows/Libwaifu.Gui/Views/SpeechTaskView.xaml b/waifu-app/wpf/Libwaifu.Gui/Views/SpeechTaskView.xaml similarity index 100% rename from windows/Libwaifu.Gui/Views/SpeechTaskView.xaml rename to waifu-app/wpf/Libwaifu.Gui/Views/SpeechTaskView.xaml diff --git a/windows/Libwaifu.Gui/Views/SpeechTaskView.xaml.cs b/waifu-app/wpf/Libwaifu.Gui/Views/SpeechTaskView.xaml.cs similarity index 100% rename from windows/Libwaifu.Gui/Views/SpeechTaskView.xaml.cs rename to waifu-app/wpf/Libwaifu.Gui/Views/SpeechTaskView.xaml.cs diff --git a/windows/Libwaifu.Gui/app.ico b/waifu-app/wpf/Libwaifu.Gui/app.ico similarity index 100% rename from windows/Libwaifu.Gui/app.ico rename to waifu-app/wpf/Libwaifu.Gui/app.ico diff --git a/windows/Libwaifu.Gui/app.manifest b/waifu-app/wpf/Libwaifu.Gui/app.manifest similarity index 100% rename from windows/Libwaifu.Gui/app.manifest rename to waifu-app/wpf/Libwaifu.Gui/app.manifest diff --git a/windows/README.md b/waifu-app/wpf/README.md similarity index 88% rename from windows/README.md rename to waifu-app/wpf/README.md index a0b5f8c..3988689 100644 --- a/windows/README.md +++ b/waifu-app/wpf/README.md @@ -1,8 +1,8 @@ # libwaifu-gui, on Windows -The same screen as `macos/`, in WPF. Laid out from the design spec (`libwaifu-gui-design.pdf`, -sections 2 to 13) and from the screenshots in it: one window, three columns, a fixed dark grey -that follows neither the system theme nor the wallpaper. +The same screen as `waifu-app/swiftui/`, in WPF. Laid out from the design spec +(`libwaifu-gui-design.pdf`, sections 2 to 13) and from the screenshots in it: one window, three +columns, a fixed dark grey that follows neither the system theme nor the wallpaper. It draws: the window calls `waifu.dll` through the C API in `docs/ffi.md`, in the app's own process, and a model that is on the disk is read onto the card by the first job that needs it. @@ -18,8 +18,8 @@ install. It is built as x64 because the library is: an AnyCPU program would be l ```powershell & "C:\Program Files\Microsoft Visual Studio\18\Community\MSBuild\Current\Bin\MSBuild.exe" ` - windows\Libwaifu.Gui\Libwaifu.Gui.csproj -p:Configuration=Release -windows\Libwaifu.Gui\bin\Release\libwaifu-gui.exe + waifu-app\wpf\Libwaifu.Gui\Libwaifu.Gui.csproj -p:Configuration=Release +waifu-app\wpf\Libwaifu.Gui\bin\Release\libwaifu-gui.exe ``` The library itself comes from this repository's own build: @@ -27,11 +27,12 @@ The library itself comes from this repository's own build: ```powershell cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release -DWITH_OPENMP=ON -DWITH_CUDA=ON ` -DCMAKE_CUDA_ARCHITECTURES=86-real -cmake --build build --target waifu-ffi # writes build\waifu.dll and build\include\waifu.h +cmake --build build --target waifu-modelmanager # build\waifu.dll, build\waifu_modelmanager.dll ``` The program looks for `waifu.dll` beside itself, then under `%LIBWAIFU_LIB_DIR%`, then in this -repository's `build\`. Copying it next to the exe is what an installed copy would look like. +repository's `build\`, and for `waifu_modelmanager.dll` beside whichever it found. Copying both +next to the exe is what an installed copy would look like. `-m ` chooses a model and turns to the page of the task it is for, `-device ` picks the device, and `-output ` changes where results are written for that run only. @@ -78,8 +79,9 @@ picture made from the seed. The window is given whichever one `Library.Available nothing above that line knows which it got. Around the engine, everything the library can answer comes from it: `Engine/NativeModels.cs` -reads `waifu_modelmanager_catalog_json`, `_describe_json`, `_directory`, `waifu_machine_json`, -and runs downloads through `waifu_modelmanager_fetch_async`. `Engine/Models.cs` is the one door +reads `waifu_modelmanager_catalog_json` and `_directory` from `waifu_modelmanager.dll`, and +`waifu_describe_json` and `waifu_machine_json` from `waifu.dll`, and runs downloads through +`waifu_modelmanager_fetch_async`. `Engine/Models.cs` is the one door the view models knock on, and it is what decides between the library and the stand-in. Three rules shape the interop, and they are the library's rather than this program's: @@ -124,4 +126,4 @@ pressed, and what the `cuda (offload)` device is there for. lines are read off Windows here and the library's answer is laid over them -- which is where the video memory and the accelerator list come from. - No manifests are shipped beside the program yet. `Library.cs` hands over a `manifests\` folder - if it finds one; `cargo run --features hub --example fetch_manifests -- ` is what fills it. + if it finds one; `cargo run -p waifu-tools --bin fetch_manifests -- ` is what fills it. diff --git a/waifu-cli/Cargo.toml b/waifu-cli/Cargo.toml new file mode 100644 index 0000000..9cabfca --- /dev/null +++ b/waifu-cli/Cargo.toml @@ -0,0 +1,38 @@ +[package] +name = "waifu-cli" +version = "0.1.0" +edition = "2021" +license = "MIT" +description = "The waifu command line tool, over the C APIs of libwaifu and libwaifu_modelmanager" + +# Both libraries are linked into the binary statically, and it calls nothing of theirs but their +# C APIs: `waifu::c_api` and `waifu_modelmanager::c_api`, the functions waifu.h and +# waifu_modelmanager.h declare. One file to copy, and nothing to find at run time. +[[bin]] +name = "waifu" +path = "src/main.rs" + +[dependencies] +# The engine and the model manager, for their C APIs only. +waifu = { path = "../waifu" } +waifu-modelmanager = { path = "../waifu-modelmanager" } + +# The terminal the whole of it happens in: which task, which model -- fetched there, under a bar +# -- and which device, and then the screen a picture is drawn from. Trimmed to the crossterm +# backend: the default set also builds the termwiz and termina ones, which is two more terminal +# stacks than a program that runs in one terminal needs. +ratatui = { version = "0.29", default-features = false, features = [ + "crossterm", + "underline-color", +] } + +# What the libraries say about models, devices and where models are kept, which they say as JSON. +serde_json = "1" + +# Reading the picture a run starts from, and writing the one it ends with: the two formats anyone +# actually hands a diffusion model, which are also the two written in pure Rust. The engine takes +# and hands back rows of RGB; the files are this tool's. +# +# Cargo.lock holds 0.25.9 rather than the newest 0.25: 0.25.10 asks for a rustc a minor version +# ahead of the one this is built with. +image = { version = "0.25", default-features = false, features = ["png", "jpeg"] } diff --git a/waifu-cli/build.rs b/waifu-cli/build.rs new file mode 100644 index 0000000..f4c1a13 --- /dev/null +++ b/waifu-cli/build.rs @@ -0,0 +1,90 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! Links what only a final link can be given, and stamps the commit the tool was built from. +//! +//! The waifu crate's build script names flint's static libraries, and those reach this link on +//! their own. Its link arguments do not: cargo keeps a `rustc-link-arg` to the package that said +//! it, so whatever flint's link flags carry as one is said again here, for this binary's link. + +use std::path::{Path, PathBuf}; + +fn main() { + let manifest = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + let lib_dir = std::env::var("LIBWAIFU_LIB_DIR") + .map(PathBuf::from) + .unwrap_or_else(|_| manifest.join("../build")); + let flags_path = lib_dir.join("flint_link_flags.txt"); + let flags = std::fs::read_to_string(&flags_path).unwrap_or_else(|error| { + panic!( + "cannot read {}: {error}\nBuild with CMake first, or point LIBWAIFU_LIB_DIR at a \ + directory that has one.", + flags_path.display() + ) + }); + println!("cargo:rerun-if-env-changed=LIBWAIFU_LIB_DIR"); + println!("cargo:rerun-if-changed={}", flags_path.display()); + + for line in flags.lines().filter(|line| line.starts_with("cargo:rustc-link-arg=")) { + println!("{line}"); + } + + stamp_revision(); +} + +/// Put the commit this was built from into the binary, for the screen to show. +/// +/// A screenshot of a run is worth little without it: two builds of the same version can differ by +/// anything, and the first thing to ask about a picture that came out wrong is which code made it. +/// Missing on purpose is a failure: a copy of the source with no git around it still builds, and +/// says "unknown" where the hash would be rather than refusing. +fn stamp_revision() { + let manifest = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + let git = manifest.join("../.git"); + + // Rebuilt when the checkout moves. HEAD covers a checkout, and the file HEAD points at covers + // a commit made on the branch already checked out. + if git.is_dir() { + println!("cargo:rerun-if-changed={}", git.join("HEAD").display()); + if let Some(reference) = head_reference(&git) { + println!("cargo:rerun-if-changed={}", git.join(reference).display()); + } + } + + let revision = std::process::Command::new("git") + .args(["rev-parse", "--short", "HEAD"]) + .current_dir(&manifest) + .output() + .ok() + .filter(|out| out.status.success()) + .and_then(|out| String::from_utf8(out.stdout).ok()) + .map(|text| text.trim().to_string()) + .filter(|text| !text.is_empty()) + .unwrap_or_else(|| "unknown".to_string()); + + println!("cargo:rustc-env=WAIFU_REVISION={revision}"); +} + +/// The ref file HEAD names, relative to the git directory, when it names one at all. A detached +/// HEAD holds a commit instead, and then there is no second file to watch. +fn head_reference(git: &Path) -> Option { + let head = std::fs::read_to_string(git.join("HEAD")).ok()?; + let reference = head.trim().strip_prefix("ref: ")?; + Some(reference.to_string()) +} diff --git a/waifu/src/cli/args.rs b/waifu-cli/src/args.rs similarity index 91% rename from waifu/src/cli/args.rs rename to waifu-cli/src/args.rs index 47cc43b..3e693ac 100644 --- a/waifu/src/cli/args.rs +++ b/waifu-cli/src/args.rs @@ -21,9 +21,9 @@ use std::fmt; -use crate::cli::task::Task; +use crate::task::Task; #[allow(unused_imports)] -use crate::{Device, Residency}; +use crate::runtime::{Device, Residency}; pub use crate::runtime::{DeviceOption, Room, Runtime}; @@ -196,9 +196,9 @@ pub fn print_options() { \".yaml\" and names the packages the weights are in, or the name of a published model, \ which is fetched before the drawing screen opens. Left out, the terminal asks for the \ task and the model, and fetches what is picked. The names are: {}.", - crate::cli::hub::names() + crate::manager::pictures() .into_iter() - .filter(|name| !crate::cli::hub::is_voice(name)) + .map(|model| model.name) .collect::>() .join(", ") ); @@ -226,18 +226,12 @@ mod tests { Args::parse(&arguments.iter().map(|a| a.to_string()).collect::>()) } - const GB: u64 = 1_000_000_000; - #[test] - fn a_named_device_is_not_second_guessed_by_the_size_of_the_model() { - assert_eq!( - DeviceOption::Cuda.resolve_for(Some(1000 * GB)), - Runtime::on(Device::Cuda) - ); - assert_eq!( - DeviceOption::Cpu.resolve_for(Some(1)), - Runtime::on(Device::Cpu) - ); + fn a_named_device_is_not_second_guessed_by_what_the_machine_has() { + let cpu_only = [Runtime::on(Device::Cpu)]; + assert_eq!(DeviceOption::Cuda.resolve(&cpu_only), Runtime::on(Device::Cuda)); + assert_eq!(DeviceOption::Auto.resolve(&cpu_only), Runtime::on(Device::Cpu)); + assert_eq!(DeviceOption::Auto.resolve(&Runtime::ALL), Runtime::on(Device::Cuda)); } #[test] @@ -257,25 +251,25 @@ mod tests { DeviceOption::Cuda, DeviceOption::Metal, ] { - assert_eq!(device.resolve().residency(), Residency::Device); + assert_eq!(device.resolve(&Runtime::ALL).residency(), Residency::Device); } let asked = args(&["-device", "metal_disk_offload"]).unwrap(); - assert_eq!(asked.device().unwrap().resolve(), Runtime::METAL_DISK_OFFLOAD); + assert_eq!(asked.device().unwrap().resolve(&Runtime::ALL), Runtime::METAL_DISK_OFFLOAD); assert_eq!( - asked.device().unwrap().resolve().residency(), + asked.device().unwrap().resolve(&Runtime::ALL).residency(), Residency::LowVram ); - assert_eq!(asked.device().unwrap().resolve().device(), Device::Metal); + assert_eq!(asked.device().unwrap().resolve(&Runtime::ALL).device(), Device::Metal); let asked = args(&["-device", "cuda_cpu_offload"]).unwrap(); - assert_eq!(asked.device().unwrap().resolve(), Runtime::CUDA_CPU_OFFLOAD); + assert_eq!(asked.device().unwrap().resolve(&Runtime::ALL), Runtime::CUDA_CPU_OFFLOAD); assert_eq!( - asked.device().unwrap().resolve().residency(), + asked.device().unwrap().resolve(&Runtime::ALL).residency(), Residency::LowVram ); assert_eq!( - asked.device().unwrap().resolve().device(), + asked.device().unwrap().resolve(&Runtime::ALL).device(), Device::Cuda, "the weights can only wait off a card there is a card for" ); @@ -308,7 +302,7 @@ mod tests { for runtime in Runtime::ALL { let named = args(&["-device", runtime.name()]).unwrap(); assert_eq!( - named.device().unwrap().resolve(), + named.device().unwrap().resolve(&Runtime::ALL), runtime, "{}", runtime.name() @@ -377,14 +371,14 @@ mod tests { #[test] fn a_named_device_is_the_one_that_is_used() { - assert_eq!(DeviceOption::Cpu.resolve().device(), Device::Cpu); - assert_eq!(DeviceOption::Cuda.resolve().device(), Device::Cuda); + assert_eq!(DeviceOption::Cpu.resolve(&Runtime::ALL).device(), Device::Cpu); + assert_eq!(DeviceOption::Cuda.resolve(&Runtime::ALL).device(), Device::Cuda); assert_eq!( - DeviceOption::CudaCpuOffload.resolve().device(), + DeviceOption::CudaCpuOffload.resolve(&Runtime::ALL).device(), Device::Cuda ); assert_eq!( - DeviceOption::MetalDiskOffload.resolve().device(), + DeviceOption::MetalDiskOffload.resolve(&Runtime::ALL).device(), Device::Metal ); } diff --git a/waifu-cli/src/c.rs b/waifu-cli/src/c.rs new file mode 100644 index 0000000..02cad45 --- /dev/null +++ b/waifu-cli/src/c.rs @@ -0,0 +1,75 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! The two libraries' C APIs, which are all of them this tool calls. +//! +//! Both are linked into the binary statically, so these are the very functions waifu.h and +//! waifu_modelmanager.h declare, with the types they declare, called as a C program would call +//! them. Nothing else of either crate is reached from here or from anywhere in the tool. + +#![allow(unused_imports)] + +use std::ffi::{c_char, CStr, CString}; + +pub use waifu::c_api::*; +pub use waifu_modelmanager::c_api::*; + +/// A string the library handed back for the caller to read, or None for NULL. +pub unsafe fn text(pointer: *const c_char) -> Option { + (!pointer.is_null()).then(|| CStr::from_ptr(pointer).to_string_lossy().into_owned()) +} + +/// `text` as a C string, with any NUL in it -- which C would read as the end -- made a space. +pub fn c_string(text: &str) -> CString { + CString::new(text.replace('\0', " ")).unwrap_or_default() +} + +/// This thread's last message from libwaifu. +pub fn waifu_error() -> String { + unsafe { text(waifu_last_error()) }.unwrap_or_default() +} + +/// This thread's last message from libwaifu_modelmanager. +pub fn manager_error() -> String { + unsafe { text(waifu_modelmanager_last_error()) }.unwrap_or_default() +} + +/// A string libwaifu returned, read and freed. +pub fn waifu_string(said: *mut c_char) -> Option { + let read = unsafe { text(said) }; + unsafe { waifu_string_free(said) }; + read +} + +/// A string libwaifu_modelmanager returned, read and freed. +pub fn manager_string(said: *mut c_char) -> Option { + let read = unsafe { text(said) }; + unsafe { waifu_modelmanager_string_free(said) }; + read +} + +/// JSON libwaifu returned, read and freed. +pub fn waifu_json(said: *mut c_char) -> Option { + waifu_string(said).and_then(|said| serde_json::from_str(&said).ok()) +} + +/// JSON libwaifu_modelmanager returned, read and freed. +pub fn manager_json(said: *mut c_char) -> Option { + manager_string(said).and_then(|said| serde_json::from_str(&said).ok()) +} diff --git a/waifu/src/cli/draw.rs b/waifu-cli/src/draw.rs similarity index 60% rename from waifu/src/cli/draw.rs rename to waifu-cli/src/draw.rs index 25b0afa..50426da 100644 --- a/waifu/src/cli/draw.rs +++ b/waifu-cli/src/draw.rs @@ -27,11 +27,11 @@ use std::path::{Path, PathBuf}; -use crate::cli::args::Args; -use crate::cli::hub; -use crate::cli::task::{Launch, Task}; -use crate::cli::tui; -use crate::describe; +use crate::args::Args; +use crate::engine::{self, Chosen}; +use crate::manager::{self, Fetch, Word}; +use crate::task::{Launch, Task}; +use crate::tui; type Error = Box; @@ -39,7 +39,7 @@ fn print_usage() { eprintln!("Usage: waifu draw [OPTIONS]"); eprintln!(); eprintln!("Options:"); - crate::cli::args::print_options(); + crate::args::print_options(); eprintln!(); } @@ -105,9 +105,6 @@ pub fn main(arguments: &[String]) -> Result<(), Error> { std::fs::create_dir_all(&output) .map_err(|error| format!("could not make {}: {error}", output.display()))?; - // Before the first screen goes up, and it stays for the rest of the run. - tui::give_the_screen_back_if_flint_aborts(); - // What this machine can run on and how much room its cards have, asked once and before any // screen is up: the first question put to a device has the tensor library say what it found. let machine = tui::Machine::measure(device); @@ -117,7 +114,14 @@ pub fn main(arguments: &[String]) -> Result<(), Error> { // that rewrites itself rather than under a bar. Some(model) => { let task = task.unwrap_or(Task::Txt2Img); - with_usage(fits(task, &model))?; + // What can be told before anything is fetched: a voice is a voice by the catalogue, + // and a manifest already here says what it can do. + if manager::is_a_voice(&model) { + return with_usage(Err(is_a_voice(&model).into())); + } + if let Ok(manifest) = manager::manifest_path(&model) { + with_usage(fits(task, &model, &engine::look_at(&manifest)?))?; + } (task, model) } // Not named: the terminal asks, and fetches what is picked under its own bar. @@ -130,27 +134,44 @@ pub fn main(arguments: &[String]) -> Result<(), Error> { // After the fetch, since it is the packages on the disk that are measured. let manifest = manifest_of(&model)?; - let (runtime, why) = machine.start_for(&model); + let chosen = engine::look_at(&manifest)?; + with_usage(fits(task, &model, &chosen))?; + let (runtime, why) = machine.start_for(chosen.bytes); let launch = Launch { task, - model, + model: manifest.to_string_lossy().into_owned(), + label: label_of(&model), runtime, }; - tui::draw(&launch, why, &manifest, image, &output, &machine) + tui::draw(&launch, why, &chosen, image, &output, &machine) } -/// Whether `model` can do `task`, as far as can be told before reading it: a published voice or -/// converter cannot draw, and a picture model can start from a picture or not by its kind. A -/// manifest on the disk is taken at its word, and says so when it is read if not. -fn fits(task: Task, model: &str) -> Result<(), String> { - if hub::is_voice(model) || hub::is_conversion(model) { - return Err(format!( - "{model} is a voice, and the terminal draws pictures: speech is the engine's, for the \ - programs built over libwaifu" - )); +/// What the screens call a model: as typed when it was a name, because that is what someone +/// would type again, and the file name when it was a path, because the whole path does not fit +/// and its tail is the part that identifies it. +pub fn label_of(model: &str) -> String { + match model.rsplit_once(['/', '\\']) { + Some((_, file)) if !file.is_empty() => file.to_string(), + _ => model.to_string(), + } +} + +fn is_a_voice(model: &str) -> String { + format!( + "{model} is a voice, and the terminal draws pictures: speech is the engine's, for the \ + programs built over libwaifu" + ) +} + +/// Whether `model` can do `task`, as its manifest says before any of its weights are read: a +/// voice or a converter cannot draw, and a picture model can start from a picture or not by its +/// kind. +fn fits(task: Task, model: &str, chosen: &Chosen) -> Result<(), String> { + if !chosen.draws { + return Err(is_a_voice(model)); } if task == Task::Img2Img { - if let Some(why) = describe::look_at(model).no_picture_because { + if let Some(why) = &chosen.no_picture_because { return Err(format!("{model} cannot do img2img: {why}")); } } @@ -158,104 +179,126 @@ fn fits(task: Task, model: &str) -> Result<(), String> { Ok(()) } -/// Whether a picture model can start from a picture, as far as is known without reading it: what -/// the terminal's img2img list is filtered on. +/// Whether a picture model can start from a picture, as far as is known without fetching it: +/// what the terminal's img2img list is filtered on. Asked of its manifest where one is here -- +/// fetched, or bundled -- and taken as yes where none is, which a run says otherwise of later. pub fn draws_from_a_picture(model: &str) -> bool { - describe::look_at(model).no_picture_because.is_none() + match manager::manifest_path(model) { + Ok(manifest) => engine::look_at(&manifest).is_ok_and(|chosen| chosen.no_picture_because.is_none()), + Err(_) => true, + } } /// The manifest of `model`, fetching the model first where it is a published name that is not on /// the disk yet -- saying how it is going on one line that rewrites itself, the command line's /// version of the model list's bar. /// -/// Nothing is fetched for a model already here, which is asked of the cache first: fetching one -/// that is here would still go looking for which hub to ask, and a machine that is offline would -/// wait on that for nothing. +/// Nothing is fetched for a model already here, which is asked of the catalogue first: fetching +/// one that is here would still go looking for which hub to ask, and a machine that is offline +/// would wait on that for nothing. fn manifest_of(model: &str) -> Result { - if let Some(manifest) = describe::on_disk(model) { - return Ok(manifest); + let here = Path::new(model).is_file() + || manager::catalog() + .iter() + .any(|listed| listed.name == model && listed.cached); + if here { + return Ok(manager::manifest_path(model)?); } - let mut from = String::new(); - let mut report = |progress: hub::Progress| { - let (file, done, total, part, parts) = match progress { - hub::Progress::From { hub } => { - from = format!(" from {hub}"); - return; - } - hub::Progress::Fetching { - file, - done, - total, - part, - parts, - } => (file, done, total, part, parts), - hub::Progress::Fetched { - file, - bytes, - part, - parts, - } => (file, bytes, Some(bytes), part, parts), + let fetch = Fetch::start(model)?; + let ended = loop { + let Ok(word) = fetch.news.recv() else { + break Word::Failed("the fetch went away".to_string()); + }; + let Word::Fetching { + hub, + file, + done, + total, + part, + parts, + } = word + else { + break word; }; + if file.is_empty() { + continue; + } + let from = hub.map(|hub| format!(" from {hub}")).unwrap_or_default(); let of = match parts { 0 => String::new(), parts => format!(" ({part} of {parts})"), }; let how_far = match total { - Some(total) if total > 0 => format!( + Some(total) => format!( "{:.0}% of {:.2} GB", done as f64 / total as f64 * 100.0, total as f64 / 1e9 ), - _ => format!("{:.2} GB", done as f64 / 1e9), + None => format!("{:.2} GB", done as f64 / 1e9), }; eprint!("\r\x1b[2Kfetching {file}{from}{of}: {how_far}"); }; - - let fetched = hub::resolve_reporting(model, &mut report, &|| false); eprintln!(); - fetched + + match ended { + Word::Done => Ok(manager::manifest_path(model)?), + Word::Stopped(said) | Word::Failed(said) => Err(said.into()), + Word::Fetching { .. } => unreachable!("the loop ends on anything else"), + } } #[cfg(test)] mod tests { use super::*; + fn a_manifest(stem: &str, kind: &str) -> PathBuf { + let path = std::env::temp_dir().join(format!("waifu-draw-cmd-{}-{stem}.yaml", std::process::id())); + std::fs::write( + &path, + format!("weights:\n - w.safetensors\nconfig:\n model:\n type: {kind}\n"), + ) + .unwrap(); + path + } + #[test] fn a_voice_is_not_something_to_draw_with() { - for voice in ["indextts", "cosyvoice"] { - let error = fits(Task::Txt2Img, voice).unwrap_err(); - assert!(error.contains("is a voice"), "{error}"); - } - assert!(fits(Task::Txt2Img, "sdxl:base").is_ok()); + let voice = engine::look_at(&a_manifest("voice", "indextts2")).unwrap(); + let error = fits(Task::Txt2Img, "indextts", &voice).unwrap_err(); + assert!(error.contains("is a voice"), "{error}"); + + let sdxl = engine::look_at(&a_manifest("sdxl", "sdxl")).unwrap(); + assert!(fits(Task::Txt2Img, "sdxl:base", &sdxl).is_ok()); } #[test] fn img2img_is_refused_for_a_model_that_cannot_start_from_a_picture() { // Said of the kind, before a package has been looked at: Anima's layer that reads a // picture is not written. - let error = fits(Task::Img2Img, "anima:turbo").unwrap_err(); + let anima = engine::look_at(&a_manifest("anima", "anima")).unwrap(); + let error = fits(Task::Img2Img, "anima:turbo", &anima).unwrap_err(); assert!(error.contains("cannot do img2img"), "{error}"); - assert!(!draws_from_a_picture("anima:turbo")); - assert!(fits(Task::Img2Img, "sdxl:base").is_ok()); - assert!(draws_from_a_picture("sdxl:base")); + let sdxl = engine::look_at(&a_manifest("sdxl2", "sdxl")).unwrap(); + assert!(fits(Task::Img2Img, "sdxl:base", &sdxl).is_ok()); // And txt2img asks nothing of it. - assert!(fits(Task::Txt2Img, "anima:turbo").is_ok()); + assert!(fits(Task::Txt2Img, "anima:turbo", &anima).is_ok()); } #[test] fn a_manifest_named_by_its_path_is_its_own_manifest() { - let path = std::env::temp_dir().join(format!("waifu-draw-cmd-{}.yaml", std::process::id())); - std::fs::write( - &path, - "weights:\n - w.safetensors\nconfig:\n model:\n type: sdxl\n", - ) - .unwrap(); - + let path = a_manifest("own", "sdxl"); // Nothing to fetch, so nothing goes looking for a hub. assert_eq!(manifest_of(path.to_str().unwrap()).unwrap(), path); + assert!(draws_from_a_picture(path.to_str().unwrap())); let _ = std::fs::remove_file(&path); } + + #[test] + fn a_model_is_called_by_its_name_or_its_file() { + assert_eq!(label_of("sdxl:base"), "sdxl:base"); + assert_eq!(label_of("/models/sdxl-base.yaml"), "sdxl-base.yaml"); + } } diff --git a/waifu-cli/src/engine.rs b/waifu-cli/src/engine.rs new file mode 100644 index 0000000..70edd7a --- /dev/null +++ b/waifu-cli/src/engine.rs @@ -0,0 +1,396 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! libwaifu, as the drawing screen asks it: what a model is, and an engine on a device that reads +//! it and draws with it. +//! +//! The types are the screen's own -- the boxes it fills in, the progress it draws a bar from -- +//! and what they are filled from is what libwaifu says through its C API. + +use std::ffi::c_void; +use std::path::Path; +use std::ptr; +use std::sync::mpsc::{channel, Sender}; + +use crate::c::{ + self, WaifuCompletionEvent, WaifuDrawRequest, WaifuEngine, WaifuImage, + WaifuImageCompletionEvent, WaifuImageInput, WaifuJob, WaifuProgressEvent, WaifuStage, + WaifuStatusCode, +}; +use crate::runtime::Runtime; + +/// The sizes offered for a model whose manifest does not name its own: the ones SDXL was trained +/// at, which is also what libwaifu offers for one. +pub const SIZES: &[(i32, i32)] = &[ + (512, 512), + (640, 640), + (768, 768), + (896, 896), + (1024, 1024), + (832, 1216), + (1216, 832), + (768, 1344), + (1344, 768), +]; + +/// What a run is asked for. +#[derive(Clone, Debug, PartialEq)] +pub struct GenerationOptions { + pub width: i32, + pub height: i32, + pub num_steps: i32, + pub guidance_scale: f32, + pub negative_prompt: String, + /// None picks one; the screen always names one, so that the picture can say which. + pub seed: Option, + /// With a picture to start from: how far the run walks away from it, 0 to 1. + pub strength: f32, +} + +impl Default for GenerationOptions { + fn default() -> GenerationOptions { + GenerationOptions { + width: 1024, + height: 1024, + num_steps: 30, + guidance_scale: 5.0, + negative_prompt: String::new(), + seed: None, + strength: 0.8, + } + } +} + +/// What the boxes start at for a model. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct GenerationDefaults { + pub width: i32, + pub height: i32, + pub num_steps: i32, + pub guidance_scale: f32, +} + +impl Default for GenerationDefaults { + fn default() -> GenerationDefaults { + let options = GenerationOptions::default(); + GenerationDefaults { + width: options.width, + height: options.height, + num_steps: options.num_steps, + guidance_scale: options.guidance_scale, + } + } +} + +/// How far along a run is. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum GenerationProgress { + Encoding, + Step { done: i32, total: i32 }, + Decoding, +} + +/// The model the screen is set to draw with, as libwaifu describes it from its manifest before +/// any of its weights are read. +#[derive(Clone, Debug, PartialEq)] +pub struct Chosen { + /// Whether it draws: a voice does not, and the terminal draws pictures. + pub draws: bool, + pub defaults: GenerationDefaults, + pub sizes: Vec<(i32, i32)>, + pub takes_guidance: bool, + pub suggested_prompt: Option, + pub suggested_avoid: Option, + /// Why it cannot be handed a picture to start from, or None where it can. + pub no_picture_because: Option, + /// How many bytes its files are, where they are all on the disk. + pub bytes: Option, +} + +impl Default for Chosen { + fn default() -> Chosen { + Chosen { + draws: true, + defaults: GenerationDefaults::default(), + sizes: SIZES.to_vec(), + takes_guidance: true, + suggested_prompt: None, + suggested_avoid: None, + no_picture_because: None, + bytes: None, + } + } +} + +/// What libwaifu says of the manifest at `manifest`, or why it says nothing. +/// What libwaifu says the model can be asked for, and what libwaifu_modelmanager says its boxes +/// start at, as one. +pub fn look_at(manifest: &Path) -> Result { + let asked = c::c_string(&manifest.to_string_lossy()); + let mut said = c::waifu_json(unsafe { c::waifu_describe_json(asked.as_ptr()) }).ok_or_else(c::waifu_error)?; + let defaults = c::manager_json(unsafe { c::waifu_modelmanager_defaults_json(asked.as_ptr()) }) + .ok_or_else(c::manager_error)?; + // The two halves of one description. A `conversion` is in both, and its defaults go inside. + if let (Some(said), Some(defaults)) = (said.as_object_mut(), defaults.as_object()) { + for (key, value) in defaults { + match (said.get_mut(key), value) { + (Some(serde_json::Value::Object(inner)), serde_json::Value::Object(more)) => { + inner.extend(more.clone()) + } + _ => { + said.insert(key.clone(), value.clone()); + } + } + } + } + Ok(chosen_from(&said)) +} + +/// A description, as `waifu_describe_json` and `waifu_modelmanager_defaults_json` write one +/// between them. +pub fn chosen_from(said: &serde_json::Value) -> Chosen { + let defaults = GenerationDefaults::default(); + let sizes: Vec<(i32, i32)> = said["sizes"] + .as_array() + .into_iter() + .flatten() + .filter_map(|size| Some((size[0].as_i64()? as i32, size[1].as_i64()? as i32))) + .collect(); + Chosen { + draws: said["kind"] == "image", + defaults: GenerationDefaults { + width: said["width"].as_i64().map_or(defaults.width, |width| width as i32), + height: said["height"].as_i64().map_or(defaults.height, |height| height as i32), + num_steps: said["steps"].as_i64().map_or(defaults.num_steps, |steps| steps as i32), + guidance_scale: said["guidance"] + .as_f64() + .map_or(defaults.guidance_scale, |guidance| guidance as f32), + }, + sizes: match sizes.is_empty() { + true => SIZES.to_vec(), + false => sizes, + }, + takes_guidance: said["takes_guidance"].as_bool().unwrap_or(true), + suggested_prompt: said["prompt"].as_str().map(str::to_string), + suggested_avoid: said["avoid"].as_str().map(str::to_string), + no_picture_because: said["why_not"]["start_from"].as_str().map(str::to_string), + bytes: said["bytes"].as_u64(), + } +} + +/// What an engine's job ends with: a picture as rows of RGB, or nothing. +#[derive(Debug, PartialEq)] +pub enum Ended { + Done(T), + Cancelled, + Failed(String), +} + +/// What a job's callbacks say, down the channel the job was started with. +#[derive(Debug, PartialEq)] +pub enum Heard { + Progress(GenerationProgress), + Loaded(Ended<()>), + Drawn(Ended<(u32, u32, Vec)>), +} + +/// libwaifu's engine on one device: a thread of its own in the library, holding at most one model. +pub struct Engine { + handle: *mut WaifuEngine, +} + +// The library's engine is called from any thread; what it holds stays on its own. +unsafe impl Send for Engine {} + +impl Engine { + /// Starts an engine on `runtime`, or says why it cannot. + pub fn new(runtime: Runtime) -> Result { + let handle = c::waifu_engine_new(runtime.waifu_device()); + match handle.is_null() { + true => Err(c::waifu_error()), + false => Ok(Engine { handle }), + } + } + + /// Reads the model whose manifest is at `manifest`, and waits for the end: how it went. + pub fn load(&self, manifest: &Path) -> Ended<()> { + let (sender, heard) = channel(); + let sender = Box::new(sender); + let model = c::c_string(&manifest.to_string_lossy()); + let job = unsafe { + c::waifu_engine_load_async( + self.handle, + model.as_ptr(), + &*sender as *const Sender as *mut c_void, + None, + Some(on_loaded), + ) + }; + if job == 0 { + return Ended::Failed(c::waifu_error()); + } + // Exactly one end comes for a job with an id; the sender lives until it has. + loop { + match heard.recv() { + Ok(Heard::Loaded(ended)) => return ended, + Ok(_) => continue, + Err(_) => return Ended::Failed("the engine went away".to_string()), + } + } + } + + /// Starts drawing; what it says comes down `sender`, ending with [`Heard::Drawn`]. `from` is + /// the picture to start from, as rows of RGB, at any size. The job, to cancel it by. + pub fn draw( + &self, + prompt: &str, + options: &GenerationOptions, + from: Option<(u32, u32, &[u8])>, + sender: &Sender, + ) -> Result { + let prompt = c::c_string(prompt); + let negative = c::c_string(&options.negative_prompt); + let key = c::c_string("start_from"); + let image = from.map(|(width, height, rgb)| WaifuImage { + width, + height, + rgb: rgb.as_ptr(), + len: rgb.len(), + }); + let input = image.as_ref().map(|image| WaifuImageInput { + key: key.as_ptr(), + image, + }); + let request = WaifuDrawRequest { + struct_size: std::mem::size_of::() as u32, + prompt: prompt.as_ptr(), + negative: negative.as_ptr(), + images: input.as_ref().map_or(ptr::null(), |input| input as *const WaifuImageInput), + image_count: input.is_some() as usize, + width: options.width.max(0) as u32, + height: options.height.max(0) as u32, + steps: options.num_steps.max(0) as u32, + guidance: options.guidance_scale, + seed: options.seed.unwrap_or_default(), + strength: options.strength, + control_scale: 1.0, + }; + let job = unsafe { + c::waifu_engine_draw_async( + self.handle, + &request, + sender as *const Sender as *mut c_void, + Some(on_progress), + Some(on_drawn), + ) + }; + match job { + 0 => Err(c::waifu_error()), + job => Ok(job), + } + } + + /// Stops a job after the step it is on. + pub fn cancel(&self, job: WaifuJob) { + unsafe { c::waifu_engine_cancel(self.handle, job) }; + } +} + +impl Drop for Engine { + /// Stops what runs, ends what waits, drops the model and joins the library's thread: no + /// callback of this engine's runs after this. + fn drop(&mut self) { + unsafe { c::waifu_engine_free(self.handle) }; + } +} + +fn status_of(code: WaifuStatusCode, message: Option) -> Ended<()> { + match code { + WaifuStatusCode::WAIFU_OK => Ended::Done(()), + WaifuStatusCode::WAIFU_CANCELLED => Ended::Cancelled, + _ => Ended::Failed(message.unwrap_or_else(|| "it failed".to_string())), + } +} + +unsafe extern "C" fn on_loaded(user_data: *mut c_void, event: *const WaifuCompletionEvent) { + let (sender, status) = (&*(user_data as *const Sender), &(*event).status); + let _ = sender.send(Heard::Loaded(status_of(status.code, c::text(status.message)))); +} + +unsafe extern "C" fn on_progress(user_data: *mut c_void, event: *const WaifuProgressEvent) { + let (sender, event) = (&*(user_data as *const Sender), &*event); + let progress = match event.stage { + WaifuStage::WAIFU_STAGE_ENCODING => GenerationProgress::Encoding, + WaifuStage::WAIFU_STAGE_DRAWING => GenerationProgress::Step { + done: event.done as i32, + total: event.total as i32, + }, + WaifuStage::WAIFU_STAGE_DECODING => GenerationProgress::Decoding, + _ => return, + }; + let _ = sender.send(Heard::Progress(progress)); +} + +unsafe extern "C" fn on_drawn(user_data: *mut c_void, event: *const WaifuImageCompletionEvent) { + let (sender, event) = (&*(user_data as *const Sender), &*event); + let ended = match (status_of(event.status.code, c::text(event.status.message)), event.image.as_ref()) { + (Ended::Done(()), Some(image)) => { + let rgb = std::slice::from_raw_parts(image.rgb, image.len).to_vec(); + Ended::Done((image.width, image.height, rgb)) + } + (Ended::Done(()), None) => Ended::Failed("the run ended with no picture".to_string()), + (Ended::Cancelled, _) => Ended::Cancelled, + (Ended::Failed(message), _) => Ended::Failed(message), + }; + let _ = sender.send(Heard::Drawn(ended)); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_description_fills_the_boxes_and_falls_back_where_it_says_nothing() { + let said = serde_json::json!({ + "kind": "image", "width": 832, "height": 1216, "steps": 8, "guidance": 1.0, + "takes_guidance": false, "sizes": [[832, 1216], [1024, 1024], [1216, 832]], + "prompt": "masterpiece", "avoid": null, + "why_not": { "start_from": "Anima cannot start from a picture yet" }, + "bytes": 6938000000u64, + }); + let chosen = chosen_from(&said); + assert!(chosen.draws); + assert_eq!(chosen.defaults.num_steps, 8); + assert_eq!(chosen.sizes.len(), 3); + assert!(!chosen.takes_guidance); + assert_eq!(chosen.suggested_prompt.as_deref(), Some("masterpiece")); + assert!(chosen.no_picture_because.unwrap().contains("Anima")); + assert_eq!(chosen.bytes, Some(6_938_000_000)); + + let voice = chosen_from(&serde_json::json!({ "kind": "speech" })); + assert!(!voice.draws); + assert_eq!(voice.sizes, SIZES.to_vec()); + } + + #[test] + fn a_manifest_that_is_not_there_is_said_why() { + let error = look_at(Path::new("/no/such/model.yaml")).unwrap_err(); + assert!(error.contains("/no/such/model.yaml"), "{error}"); + } +} diff --git a/waifu-cli/src/log.rs b/waifu-cli/src/log.rs new file mode 100644 index 0000000..e6015ab --- /dev/null +++ b/waifu-cli/src/log.rs @@ -0,0 +1,110 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! Where the libraries' log lines go: one sink for both of them, in this process. +//! +//! Each library writes its lines to a callback once it is handed one, and this hands each of them +//! one that puts every line through [`hand_over`] -- the console until a screen sets a sink of its +//! own, as the drawing screen does. A fatal line is the last thing that happens before the process +//! aborts, so the terminal is given back before it is written, wherever it goes. + +use std::ffi::c_void; +use std::ptr; +use std::sync::RwLock; + +use crate::c::{self, WaifuLogEvent, WaifuLogLevel, WaifuModelManagerLogEvent, WaifuModelManagerLogLevel}; + +/// How much a line matters. +#[derive(Clone, Copy, Debug, Eq, PartialEq, PartialOrd, Ord)] +pub enum Level { + Debug, + Info, + Warning, + Error, + /// The last line: the process aborts once the sink returns. + Fatal, +} + +impl Level { + pub fn name(self) -> &'static str { + match self { + Level::Debug => "DEBUG", + Level::Info => "INFO", + Level::Warning => "WARNING", + Level::Error => "ERROR", + Level::Fatal => "FATAL", + } + } +} + +/// What receives a line: its level, where it was written, and the line alone. Called on +/// whichever thread wrote it. +pub type Sink = Box; + +static SINK: RwLock> = RwLock::new(None); + +/// Sends every line to `sink` from now on, or back to the console for `None`. +pub fn set_sink(sink: Option) { + *SINK.write().unwrap_or_else(|held| held.into_inner()) = sink; +} + +/// Sends both libraries' lines here, from now on. Called once, first: the tensor library says +/// what hardware it found the first time a device is asked about. +pub fn listen() { + c::waifu_set_log_callback(ptr::null_mut(), Some(from_waifu)); + c::waifu_modelmanager_set_log_callback(ptr::null_mut(), Some(from_the_model_manager)); +} + +/// A line to the sink where there is one, and the console where there is not. +fn hand_over(level: Level, source: &str, message: &str) { + if level == Level::Fatal { + crate::tui::give_the_screen_back(); + } + let sink = SINK.read().unwrap_or_else(|held| held.into_inner()); + match sink.as_ref() { + Some(sink) => sink(level, source, message), + None => eprintln!("{} {source}] {message}", level.name()), + } +} + +unsafe extern "C" fn from_waifu(_: *mut c_void, event: *const WaifuLogEvent) { + let event = &*event; + let level = match event.level { + WaifuLogLevel::WAIFU_LOG_DEBUG => Level::Debug, + WaifuLogLevel::WAIFU_LOG_INFO => Level::Info, + WaifuLogLevel::WAIFU_LOG_WARNING => Level::Warning, + WaifuLogLevel::WAIFU_LOG_ERROR => Level::Error, + WaifuLogLevel::WAIFU_LOG_FATAL => Level::Fatal, + }; + let source = c::text(event.source).unwrap_or_default(); + hand_over(level, &source, &c::text(event.message).unwrap_or_default()); +} + +unsafe extern "C" fn from_the_model_manager(_: *mut c_void, event: *const WaifuModelManagerLogEvent) { + let event = &*event; + let level = match event.level { + WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_DEBUG => Level::Debug, + WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_INFO => Level::Info, + WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_WARNING => Level::Warning, + WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_ERROR => Level::Error, + WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_FATAL => Level::Fatal, + }; + let source = c::text(event.source).unwrap_or_default(); + hand_over(level, &source, &c::text(event.message).unwrap_or_default()); +} diff --git a/waifu/src/cli/mod.rs b/waifu-cli/src/main.rs similarity index 92% rename from waifu/src/cli/mod.rs rename to waifu-cli/src/main.rs index d728ef8..99f3e0c 100644 --- a/waifu/src/cli/mod.rs +++ b/waifu-cli/src/main.rs @@ -17,11 +17,12 @@ // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -//! The libwaifu command line tool: draw a picture with a model. +//! The libwaifu command line tool, `waifu`: draw a picture with a model. //! -//! Behind the `cli` feature, because it is the only thing in this crate that needs a terminal -//! stack and an image encoder. The `waifu` binary is a shim over [`run`]; everything else here is -//! its implementation. +//! A program over the two libraries' C APIs, as any other program would be: libwaifu_modelmanager +//! for which models there are and fetching them, libwaifu for drawing with one -- both linked in +//! statically, so the tool is one file. See [`c`] for what it calls, and docs/ffi.md for what +//! those calls promise. //! //! There is one thing the tool does, and the command line exists to say how to do it rather than //! which to do: `waifu` on its own is the same as `waifu draw`. Both ask in the terminal what the @@ -31,13 +32,15 @@ //! screen. mod args; +mod c; mod draw; +mod engine; +mod log; +mod manager; +mod runtime; mod task; mod tui; -// In the library since the C API needed them; named here as well, where they used to be. -pub(crate) use crate::{config, hub}; - use std::process::ExitCode; type Error = Box; @@ -93,7 +96,9 @@ fn wants_usage(arguments: &[String]) -> bool { } /// Runs the command line tool: reads the process arguments and dispatches to a subcommand. -pub fn run() -> ExitCode { +fn main() -> ExitCode { + log::listen(); + let arguments: Vec = std::env::args().skip(1).collect(); // The first argument is the command, unless it starts with a dash, in which case no command diff --git a/waifu-cli/src/manager.rs b/waifu-cli/src/manager.rs new file mode 100644 index 0000000..92fe8a1 --- /dev/null +++ b/waifu-cli/src/manager.rs @@ -0,0 +1,310 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! libwaifu_modelmanager, as the screens ask it: which models are published, which are here, +//! where they are kept, and fetching one with news of how it is going down a channel. + +use std::ffi::c_void; +use std::path::{Path, PathBuf}; +use std::ptr; +use std::sync::mpsc::{channel, Receiver, Sender}; +use std::sync::Mutex; + +use crate::c::{ + self, WaifuModelFetch, WaifuModelManagerCompletionEvent as WaifuCompletionEvent, + WaifuModelManagerProgressEvent as WaifuProgressEvent, + WaifuModelManagerStatusCode as WaifuStatusCode, +}; + +/// The settings file the model directory is written into, as the screens name it. +pub const CONFIG_FILE_NAME: &str = "config.toml"; + +/// The environment variable that decides the model directory over everything else. +pub const CACHE_ENV: &str = "WAIFU_CACHE"; + +/// One published model, as the catalogue says it. +#[derive(Clone, Debug, PartialEq)] +pub struct Listed { + pub name: &'static str, + pub full_name: &'static str, + /// "image", "speech" or "conversion"; a model that does two is listed twice. + pub kind: String, + /// Whether every package of it is here. + pub cached: bool, + /// What is on the disk of it, which is most of a model for one that was interrupted. + pub bytes: u64, + pub explicit: bool, + /// The directory it is in, where any of it is on the disk. + pub directory: Option, +} + +/// A catalogue name for the rest of the run. The names are a few dozen at most and the same each +/// time the catalogue is read, so each is kept once, and the screens can hold them as plainly as +/// they could when the catalogue was a table compiled in. +pub fn kept(name: &str) -> &'static str { + static NAMES: Mutex> = Mutex::new(Vec::new()); + let mut names = NAMES.lock().unwrap_or_else(|held| held.into_inner()); + if let Some(name) = names.iter().find(|kept| **kept == name) { + return name; + } + let name: &'static str = Box::leak(name.to_string().into_boxed_str()); + names.push(name); + name +} + +/// Every published model, in the catalogue's order. +pub fn catalog() -> Vec { + let Some(said) = c::manager_json(c::waifu_modelmanager_catalog_json()) else { + return Vec::new(); + }; + said["models"] + .as_array() + .into_iter() + .flatten() + .filter_map(|model| { + Some(Listed { + name: kept(model["name"].as_str()?), + full_name: kept(model["full_name"].as_str()?), + kind: model["kind"].as_str()?.to_string(), + cached: model["cached"].as_bool().unwrap_or(false), + bytes: model["bytes_on_disk"].as_u64().unwrap_or(0), + explicit: model["explicit"].as_bool().unwrap_or(false), + directory: model["directory"].as_str().map(PathBuf::from), + }) + }) + .collect() +} + +/// The published models that draw. +pub fn pictures() -> Vec { + catalog().into_iter().filter(|model| model.kind == "image").collect() +} + +/// Whether `name` is a published model that speaks or converts rather than draws. +pub fn is_a_voice(name: &str) -> bool { + catalog() + .iter() + .any(|model| model.name == name && model.kind != "image") +} + +/// Which of the three places that can name the model directory named it. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Said { + /// `WAIFU_CACHE`, which wins over everything: it is what a script sets for one run. + Environment, + /// `model_dir` in `config.toml`, which is what the model screen writes. + Config, + /// Nobody: the platform's cache directory. + Default, +} + +impl Said { + /// How the model screen says where the directory came from. + pub fn describe(self) -> &'static str { + match self { + Said::Environment => CACHE_ENV, + Said::Config => CONFIG_FILE_NAME, + Said::Default => "default", + } + } +} + +/// Where models are fetched into, who said so, and what else there is to know about it. +#[derive(Clone, Debug, PartialEq)] +pub struct Directory { + pub now: PathBuf, + pub said: Said, + /// Where they would go if nobody said. + pub default: Option, + /// The settings file a directory chosen here is written into. + pub config_file: Option, +} + +/// Where models are fetched into, or why that cannot be told -- a settings file that does not +/// read, most likely. +pub fn directory() -> Result { + let said = c::manager_json(c::waifu_modelmanager_directory_json()) + .ok_or_else(c::manager_error)?; + Ok(Directory { + now: PathBuf::from(said["directory"].as_str().unwrap_or_default()), + said: match said["said"].as_str() { + Some("environment") => Said::Environment, + Some("config") => Said::Config, + _ => Said::Default, + }, + default: said["default"].as_str().map(PathBuf::from), + config_file: said["config_file"].as_str().map(PathBuf::from), + }) +} + +/// Models are fetched into `directory` from now on, or the default for None. +pub fn set_directory(directory: Option<&Path>) -> Result<(), String> { + let path = directory.map(|directory| c::c_string(&directory.to_string_lossy())); + let code = unsafe { + c::waifu_modelmanager_set_directory(path.as_ref().map_or(ptr::null(), |path| path.as_ptr())) + }; + match code { + WaifuStatusCode::WAIFU_MODELMANAGER_OK => Ok(()), + _ => Err(c::manager_error()), + } +} + +/// Throws away what has been fetched of `name`. +pub fn remove(name: &str) -> Result<(), String> { + let name = c::c_string(name); + match unsafe { c::waifu_modelmanager_remove(name.as_ptr()) } { + WaifuStatusCode::WAIFU_MODELMANAGER_OK => Ok(()), + _ => Err(c::manager_error()), + } +} + +/// The manifest of `model` -- a published name, or a manifest's own path -- for the engine to be +/// handed: the fetched one, or a bundled one where nothing has been. An error where there is none +/// here. +pub fn manifest_path(model: &str) -> Result { + let model = c::c_string(model); + c::manager_string(unsafe { c::waifu_modelmanager_manifest_path(model.as_ptr()) }) + .map(PathBuf::from) + .ok_or_else(c::manager_error) +} + +/// What a fetch has to say. +#[derive(Clone, Debug, PartialEq)] +pub enum Word { + /// A file, how far into it, how long it is when that is known, which package of how many it + /// is -- `parts` being zero until the first package has named the others -- and which hub it + /// comes from once that is settled. + Fetching { + hub: Option, + file: String, + done: u64, + total: Option, + part: usize, + parts: usize, + }, + /// Every package is here. + Done, + /// It was asked to stop, and did; what is said is what was kept. + Stopped(String), + Failed(String), +} + +/// What a stop leaves behind, as the screens say it. +const STOPPED: &str = "stopped -- what has come down is kept, and fetching it again carries on from there"; + +/// A fetch running in libwaifu_modelmanager, and the channel its news comes down. +pub struct Fetch { + handle: *mut WaifuModelFetch, + /// What the callbacks are handed, owned here and let go of once they cannot be called again. + sender: *mut Sender, + pub news: Receiver, +} + +// The handle is only cancelled and freed, which the library allows from any thread. +unsafe impl Send for Fetch {} + +impl Fetch { + /// Starts fetching `name`. + pub fn start(name: &str) -> Result { + let (sender, news) = channel(); + let sender = Box::into_raw(Box::new(sender)); + let model = c::c_string(name); + let handle = unsafe { + c::waifu_modelmanager_fetch_async( + model.as_ptr(), + sender as *mut c_void, + Some(on_progress), + Some(on_complete), + ) + }; + if handle.is_null() { + // Refused at the call: no callback will ever run, so the sender is this side's again. + drop(unsafe { Box::from_raw(sender) }); + return Err(c::manager_error()); + } + Ok(Fetch { handle, sender, news }) + } + + /// Asks it to stop. What has come down is kept. + pub fn stop(&self) { + unsafe { c::waifu_modelmanager_fetch_cancel(self.handle) }; + } +} + +impl Drop for Fetch { + fn drop(&mut self) { + // Stops it if it is still going, and waits for its thread: no callback runs after this. + unsafe { c::waifu_modelmanager_fetch_free(self.handle) }; + drop(unsafe { Box::from_raw(self.sender) }); + } +} + +unsafe extern "C" fn on_progress(user_data: *mut c_void, event: *const WaifuProgressEvent) { + let (sender, event) = (&*(user_data as *const Sender), &*event); + let _ = sender.send(Word::Fetching { + hub: c::text(event.hub), + file: c::text(event.file).unwrap_or_default(), + done: event.done, + total: (event.total > 0).then_some(event.total), + part: event.part as usize, + parts: event.parts as usize, + }); +} + +unsafe extern "C" fn on_complete(user_data: *mut c_void, event: *const WaifuCompletionEvent) { + let (sender, status) = (&*(user_data as *const Sender), &(*event).status); + let _ = sender.send(match status.code { + WaifuStatusCode::WAIFU_MODELMANAGER_OK => Word::Done, + WaifuStatusCode::WAIFU_MODELMANAGER_CANCELLED => Word::Stopped(STOPPED.to_string()), + _ => Word::Failed(c::text(status.message).unwrap_or_else(|| "the fetch failed".to_string())), + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_catalogue_lists_pictures_and_voices() { + let catalog = catalog(); + assert!(catalog.iter().any(|model| model.name == "sdxl:base" && model.kind == "image")); + assert!(is_a_voice("indextts")); + assert!(!is_a_voice("sdxl:base")); + assert!(pictures().iter().all(|model| model.kind == "image")); + } + + #[test] + fn a_name_is_kept_once() { + // Two strings with the same letters at two addresses, which come back as one. + let built = String::from("sdxl:") + "base"; + assert!(std::ptr::eq(kept("sdxl:base"), kept(&built))); + } + + #[test] + fn a_fetch_of_a_model_nobody_published_fails_at_the_end() { + let fetch = Fetch::start("no-such-model:v9").unwrap(); + let end = loop { + match fetch.news.recv_timeout(std::time::Duration::from_secs(20)).unwrap() { + Word::Fetching { .. } => continue, + end => break end, + } + }; + assert!(matches!(&end, Word::Failed(message) if message.contains("no-such-model:v9")), "{end:?}"); + } +} diff --git a/waifu-cli/src/runtime.rs b/waifu-cli/src/runtime.rs new file mode 100644 index 0000000..7ec0fd4 --- /dev/null +++ b/waifu-cli/src/runtime.rs @@ -0,0 +1,377 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! Where a run goes: which device, where the weights wait, and how to choose for a model. +//! +//! What this machine can run on and how much room its cards have are libwaifu's to say, through +//! `waifu_devices_json`; which of them suits a model is this tool's, and it is all here. + +/// A device a run can be sent to, as libwaifu names them. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Device { + Cpu, + Cuda, + Metal, + Vulkan, +} + +impl Device { + pub fn name(self) -> &'static str { + match self { + Device::Cpu => "cpu", + Device::Cuda => "cuda", + Device::Metal => "metal", + Device::Vulkan => "vulkan", + } + } +} + +/// Where the weights wait between the steps that read them: on the device, or -- for a model +/// larger than it -- on the host or the disk, each moved over at the step that reads it. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Residency { + Device, + LowVram, +} + +use crate::c::{self, WaifuDevice}; + +/// Where a run goes: a device, and where the weights wait between the steps that read them. +/// +/// One name rather than two answers. Where the weights wait is a question about a card the model +/// does not fit on, so it has one device it can be asked of and one reason to ask it, and asking +/// it beside the device made a second box whose rows were struck out three times out of four. As +/// a device name -- `cuda_cpu_offload`, cuda with the weights kept on the host, and +/// `metal_disk_offload`, metal with them left on the disk -- the pair that cannot be built cannot +/// be said either, and there is one list on screen instead of two. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct Runtime { + device: Device, + residency: Residency, +} + +impl Runtime { + /// Every place a run can go, in the order the flag's own usage lists them. + pub const ALL: [Runtime; 6] = [ + Runtime::on(Device::Cpu), + Runtime::on(Device::Cuda), + Runtime::CUDA_CPU_OFFLOAD, + Runtime::on(Device::Metal), + Runtime::METAL_DISK_OFFLOAD, + Runtime::on(Device::Vulkan), + ]; + + /// Cuda, with the package held on the host and each weight moved onto the card at the + /// instruction that reads it. For a card the model does not fit on. + pub const CUDA_CPU_OFFLOAD: Runtime = Runtime { + device: Device::Cuda, + residency: Residency::LowVram, + }; + + /// Metal, with the package left in its file and each weight read off the disk into the + /// device's memory at the instruction that reads it. Metal's memory is the machine's, so this + /// is for a model larger than the machine. + pub const METAL_DISK_OFFLOAD: Runtime = Runtime { + device: Device::Metal, + residency: Residency::LowVram, + }; + + /// A device with the weights on it, which is what every name but the offload ones means. + pub const fn on(device: Device) -> Runtime { + Runtime { + device, + residency: Residency::Device, + } + } + + pub fn device(self) -> Device { + self.device + } + + #[cfg_attr(not(test), allow(dead_code))] + pub fn residency(self) -> Residency { + self.residency + } + + /// What it is called on the command line and on screen. + pub fn name(self) -> &'static str { + match (self.device, self.residency) { + (device, Residency::Device) => device.name(), + (Device::Metal, Residency::LowVram) => "metal_disk_offload", + (_, Residency::LowVram) => "cuda_cpu_offload", + } + } + + /// What the drawing screen's device list says beside the name of one this machine can use. + pub fn about(self) -> &'static str { + match (self.device, self.residency) { + (_, Residency::Device) => "ready", + // The rows that have to say more than that. They are the slow answers, and the slow + // answer picked by someone who was not told is a run that looks broken. + (Device::Metal, Residency::LowVram) => { + "ready, slower: for a model larger than the memory" + } + (_, Residency::LowVram) => "ready, slower: for a model larger than the card", + } + } +} + +/// Where the model should run, as `-device` spells it. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum DeviceOption { + /// The first accelerator there is a device for -- CUDA, then Metal, then Vulkan -- and the CPU + /// otherwise. + Auto, + Cpu, + Cuda, + CudaCpuOffload, + Metal, + MetalDiskOffload, + Vulkan, +} + +impl DeviceOption { + /// The runtime this actually means on a machine that can run on `available`, which is the + /// one question `auto` leaves open. + pub fn resolve(self, available: &[Runtime]) -> Runtime { + match self { + DeviceOption::Cpu => Runtime::on(Device::Cpu), + DeviceOption::Cuda => Runtime::on(Device::Cuda), + DeviceOption::CudaCpuOffload => Runtime::CUDA_CPU_OFFLOAD, + DeviceOption::Metal => Runtime::on(Device::Metal), + DeviceOption::MetalDiskOffload => Runtime::METAL_DISK_OFFLOAD, + DeviceOption::Vulkan => Runtime::on(Device::Vulkan), + // Never an offload one: with no model in hand there is nothing to measure against + // the card. [`Runtime::best`] is the answer once there is. + // + // At most one of CUDA and Metal is ever built, so the order between those two only + // decides which check runs first. Vulkan comes last because it runs on the same cards + // as either, through kernels of its own that are not the fastest those cards have. + DeviceOption::Auto => [Device::Cuda, Device::Metal, Device::Vulkan] + .into_iter() + .map(Runtime::on) + .find(|runtime| available.contains(runtime)) + .unwrap_or(Runtime::on(Device::Cpu)), + } + } +} + +/// How much of each card is free, where the card said. `None` is a card that was not asked or +/// did not answer, which is taken as room enough: that is what `auto` assumed before it measured. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct Room { + pub cuda: Option, + pub vulkan: Option, +} + +/// What a model of `weights` bytes wants of a card to run with its weights on it: the weights, +/// and room beside them for what a run makes -- the latents, the attention, the decoded picture. +/// +/// A guess, and a generous one. Too small and the run aborts out of memory halfway through a +/// picture; too large and a model that would have fitted runs offloaded, which is slower but +/// finishes. SDXL's 7 GB comes to 9.7, and Qwen-Image's 15 GB of fp8 to 18.6, which is what puts +/// the one on a sixteen gigabyte card and the other beside it. +pub fn wanted_on_the_card(weights: u64) -> u64 { + const GIB: u64 = 1 << 30; + weights + weights / 10 + 2 * GIB +} + +impl Runtime { + /// What a model of `weights` bytes wants of this runtime's card and what the card has free, + /// where the one is more than the other: a run there is likely to abort out of memory. + /// + /// Only a runtime that puts the weights on a card that was measured can be short. The offload + /// one keeps them on the host, Metal's memory is the machine's, and a size or a card that is + /// not known is taken as one that fits. + pub fn short_of(self, weights: Option, room: Room) -> Option<(u64, u64)> { + let free = match (self.device, self.residency) { + (Device::Cuda, Residency::Device) => room.cuda, + (Device::Vulkan, Residency::Device) => room.vulkan, + _ => None, + }?; + let wanted = wanted_on_the_card(weights?); + (wanted > free).then_some((wanted, free)) + } + + /// The fastest of `available` a model of `weights` bytes will run on, given `room`. + /// + /// CUDA where the model fits on the card, and CUDA with the weights on the host where it does + /// not: slower, but far faster than the processor, and it finishes where the other aborts. + /// Then Metal, whose memory is the machine's. Then Vulkan, which has no offload of its own and + /// so is only the answer where the model fits. The processor otherwise. A size that is not + /// known is taken as one that fits. + pub fn best(weights: Option, available: &[Runtime], room: Room) -> Runtime { + let cuda = Runtime::on(Device::Cuda); + let metal = Runtime::on(Device::Metal); + let vulkan = Runtime::on(Device::Vulkan); + + if available.contains(&cuda) { + if cuda.short_of(weights, room).is_none() { + return cuda; + } + if available.contains(&Runtime::CUDA_CPU_OFFLOAD) { + return Runtime::CUDA_CPU_OFFLOAD; + } + } + if available.contains(&metal) { + return metal; + } + if available.contains(&vulkan) && vulkan.short_of(weights, room).is_none() { + return vulkan; + } + Runtime::on(Device::Cpu) + } +} + +impl Runtime { + /// The device libwaifu's engine is started on for this. + pub fn waifu_device(self) -> WaifuDevice { + match (self.device, self.residency) { + (Device::Cpu, _) => WaifuDevice::WAIFU_DEVICE_CPU, + (Device::Cuda, Residency::Device) => WaifuDevice::WAIFU_DEVICE_CUDA, + (Device::Cuda, Residency::LowVram) => WaifuDevice::WAIFU_DEVICE_CUDA_CPU_OFFLOAD, + (Device::Metal, Residency::Device) => WaifuDevice::WAIFU_DEVICE_METAL, + (Device::Metal, Residency::LowVram) => WaifuDevice::WAIFU_DEVICE_METAL_DISK_OFFLOAD, + (Device::Vulkan, _) => WaifuDevice::WAIFU_DEVICE_VULKAN, + } + } + + /// The one called `name`, as `waifu_devices_json` and the -device flag name them. + pub fn named(name: &str) -> Option { + Runtime::ALL.into_iter().find(|runtime| runtime.name() == name) + } +} + +/// What this machine can run on and how much room its cards have, asked of libwaifu. +/// +/// The first time CUDA is asked it builds a context on the card, a few hundred megabytes of it, +/// and the tensor library says what hardware it found: so this is asked once, before any screen +/// is up. +pub fn measure() -> (Vec, Room) { + let Some(said) = c::waifu_json(c::waifu_devices_json()) else { + return (vec![Runtime::on(Device::Cpu)], Room::default()); + }; + let mut available = Vec::new(); + let mut room = Room::default(); + for row in said["devices"].as_array().into_iter().flatten() { + let Some(runtime) = row["name"].as_str().and_then(Runtime::named) else { + continue; + }; + if row["available"].as_bool() != Some(true) { + continue; + } + available.push(runtime); + let free = row["free"].as_u64(); + match runtime.device() { + Device::Cuda if room.cuda.is_none() => room.cuda = free, + Device::Vulkan => room.vulkan = free, + _ => {} + } + } + if available.is_empty() { + available.push(Runtime::on(Device::Cpu)); + } + (available, room) +} + +#[cfg(test)] +mod tests { + use super::*; + + const GB: u64 = 1_000_000_000; + + #[test] + fn auto_puts_a_model_on_the_card_it_fits_and_offloads_one_it_does_not() { + let everything = Runtime::ALL.to_vec(); + let cuda = Runtime::on(Device::Cuda); + let sixteen = Room { + cuda: Some(16 * GB), + vulkan: Some(16 * GB), + }; + + // SDXL on a sixteen gigabyte card, and Qwen-Image's fp8 beside it. + assert_eq!(Runtime::best(Some(7 * GB), &everything, sixteen), cuda); + assert_eq!( + Runtime::best(Some(15 * GB), &everything, sixteen), + Runtime::CUDA_CPU_OFFLOAD + ); + + // A size or a card that did not say is the old answer, which is the card. + assert_eq!(Runtime::best(None, &everything, sixteen), cuda); + assert_eq!( + Runtime::best(Some(40 * GB), &everything, Room::default()), + cuda + ); + } + + #[test] + fn auto_without_cuda_goes_to_vulkan_only_where_the_model_fits() { + let cpu = Runtime::on(Device::Cpu); + let vulkan = Runtime::on(Device::Vulkan); + let here = [cpu, vulkan]; + let eight = Room { + cuda: None, + vulkan: Some(8 * GB), + }; + + assert_eq!(Runtime::best(Some(3 * GB), &here, eight), vulkan); + assert_eq!(Runtime::best(Some(7 * GB), &here, eight), cpu); + + // Metal has the machine's memory, so it has no card to be too big for -- and its offload + // is never chosen for anyone, however large the model: it is asked for by name. + let mac = [cpu, Runtime::on(Device::Metal), Runtime::METAL_DISK_OFFLOAD]; + assert_eq!( + Runtime::best(Some(400 * GB), &mac, Room::default()), + Runtime::on(Device::Metal) + ); + assert_eq!( + Runtime::METAL_DISK_OFFLOAD.short_of(Some(400 * GB), Room::default()), + None + ); + assert_eq!(Runtime::best(Some(1), &[cpu], Room::default()), cpu); + } + + #[test] + fn only_a_card_the_weights_go_on_can_be_short_of_room() { + let eight = Room { + cuda: Some(8 * GB), + vulkan: Some(8 * GB), + }; + let cuda = Runtime::on(Device::Cuda); + + assert_eq!(cuda.short_of(Some(3 * GB), eight), None); + assert_eq!( + cuda.short_of(Some(7 * GB), eight), + Some((wanted_on_the_card(7 * GB), 8 * GB)) + ); + assert!(Runtime::on(Device::Vulkan) + .short_of(Some(7 * GB), eight) + .is_some()); + + // The offload one keeps the weights on the host, and the processor has no card at all. + assert_eq!( + Runtime::CUDA_CPU_OFFLOAD.short_of(Some(7 * GB), eight), + None + ); + assert_eq!(Runtime::on(Device::Cpu).short_of(Some(7 * GB), eight), None); + // Nor is anything short of what was never measured. + assert_eq!(cuda.short_of(None, eight), None); + assert_eq!(cuda.short_of(Some(7 * GB), Room::default()), None); + } +} diff --git a/waifu/src/cli/task.rs b/waifu-cli/src/task.rs similarity index 94% rename from waifu/src/cli/task.rs rename to waifu-cli/src/task.rs index fce0204..6c8e647 100644 --- a/waifu/src/cli/task.rs +++ b/waifu-cli/src/task.rs @@ -25,7 +25,7 @@ //! can be watched and stopped. The device it does change, from its device row: that is a read of //! the model onto another card, which happens behind the screen. -use crate::cli::args::Runtime; +use crate::args::Runtime; /// One kind of session: a row in the terminal's first list, and the boxes the drawing screen has. #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -79,8 +79,10 @@ impl Task { #[derive(Clone, Debug, Eq, PartialEq)] pub struct Launch { pub task: Task, - /// A catalogue name or a manifest's path: what the painter hands to the hub to read. + /// The path of the model's manifest: what libwaifu is handed to read. pub model: String, + /// What the screens call it: the catalogue name it was picked by, or the manifest's file name. + pub label: String, pub runtime: Runtime, } diff --git a/waifu/src/cli/tui/ask.rs b/waifu-cli/src/tui/ask.rs similarity index 99% rename from waifu/src/cli/tui/ask.rs rename to waifu-cli/src/tui/ask.rs index 5883d16..ec7987d 100644 --- a/waifu/src/cli/tui/ask.rs +++ b/waifu-cli/src/tui/ask.rs @@ -37,9 +37,9 @@ use ratatui::text::{Line, Span}; use ratatui::widgets::{Block, BorderType, Clear, Paragraph}; use ratatui::Frame; -use crate::cli::tui::centred; -use crate::cli::tui::field::{edit_text, TextField}; -use crate::cli::tui::files::tail; +use crate::tui::centred; +use crate::tui::field::{edit_text, TextField}; +use crate::tui::files::tail; /// How wide the box is drawn, before the room it is given decides. /// diff --git a/waifu/src/cli/tui/devices.rs b/waifu-cli/src/tui/devices.rs similarity index 98% rename from waifu/src/cli/tui/devices.rs rename to waifu-cli/src/tui/devices.rs index 43d5f0f..b0b0c0f 100644 --- a/waifu/src/cli/tui/devices.rs +++ b/waifu-cli/src/tui/devices.rs @@ -35,9 +35,9 @@ use ratatui::text::{Line, Span}; use ratatui::widgets::{Block, BorderType, Clear, Paragraph}; use ratatui::Frame; -use crate::cli::args::{Room, Runtime}; -use crate::cli::tui::ask::{Answer, Confirm}; -use crate::cli::tui::{centred, gigabytes, row_style}; +use crate::args::{Room, Runtime}; +use crate::tui::ask::{Answer, Confirm}; +use crate::tui::{centred, gigabytes, row_style}; /// How wide the box is drawn, before the room it is given decides: a name, and the longest thing /// said beside one, which is how short of memory a card is. @@ -254,8 +254,8 @@ mod tests { use ratatui::backend::TestBackend; use ratatui::crossterm::event::KeyModifiers; - use crate::cli::tui::screen_text; - use crate::Device; + use crate::tui::screen_text; + use crate::runtime::Device; fn press(code: KeyCode) -> KeyEvent { KeyEvent::new(code, KeyModifiers::NONE) diff --git a/waifu/src/cli/tui/draw.rs b/waifu-cli/src/tui/draw.rs similarity index 95% rename from waifu/src/cli/tui/draw.rs rename to waifu-cli/src/tui/draw.rs index 158f0a1..81a4fbd 100644 --- a/waifu/src/cli/tui/draw.rs +++ b/waifu-cli/src/tui/draw.rs @@ -20,9 +20,9 @@ //! The last screen: a screenful of knobs, a diffusion model behind it, and a PNG at the end. //! //! A run takes minutes rather than milliseconds, which is what the shape of this is about. The -//! model sits on a thread of its own -- a tensor never leaves the thread that made it, so the -//! model cannot be touched from the drawing loop -- and the two talk through a pair of channels: -//! a prompt goes one way, and how far along the run is comes back the other. +//! model is libwaifu's, on an engine of its own, and a thread here -- the painter -- is what talks +//! to it, so that the drawing loop never waits on it: a prompt goes one way down a pair of +//! channels, and how far along the run is comes back the other. //! //! What comes back is a file name rather than a picture. A terminal is a poor place to look at //! one, so every finished run is written out where an image viewer can open it, with the prompt @@ -33,7 +33,6 @@ //! to reach. use std::io::{self, Write}; -use std::ops::ControlFlow; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc::{channel, Receiver, Sender}; @@ -47,19 +46,17 @@ use ratatui::text::{Line, Span}; use ratatui::widgets::{Block, BorderType, Paragraph, Wrap}; use ratatui::{DefaultTerminal, Frame}; -use crate::cli::args::{Room, Runtime}; -use crate::cli::hub; -use crate::cli::task::{Launch, Task}; -use crate::cli::tui::ask::{self, Answer}; -use crate::cli::tui::devices::DevicePicker; -use crate::cli::tui::field::{edit_text, TextField}; -use crate::cli::tui::files::{FilePicker, Outcome}; -use crate::cli::tui::{bar, centred, heading, take_the_screen, Machine}; -use crate::describe::{self, Chosen}; -use crate::flint::Tensor; -use crate::image_model::ImageModel; +use crate::args::{Room, Runtime}; +use crate::engine::{ + Chosen, Ended, Engine, GenerationDefaults, GenerationOptions, GenerationProgress, Heard as News, +}; use crate::log::Level; -use crate::{from_rgb8, to_rgb8, GenerationDefaults, GenerationOptions, GenerationProgress}; +use crate::task::{Launch, Task}; +use crate::tui::ask::{self, Answer}; +use crate::tui::devices::DevicePicker; +use crate::tui::field::{edit_text, TextField}; +use crate::tui::files::{FilePicker, Outcome}; +use crate::tui::{bar, centred, heading, take_the_screen, Machine}; type Error = Box; @@ -83,23 +80,20 @@ const PICTURES: &[&str] = &["png", "jpg", "jpeg"]; /// Puts the drawing screen up for `launch` and reads its model onto its device behind it, until /// somebody leaves. /// -/// `manifest` is the model's manifest, already on the disk. `why` is what `auto` said about the -/// device it picked, if it picked one; `machine` is what the device row offers instead. `image` -/// fills the draw from box, where `-i` named one, and `output` is the directory every picture is -/// written into. +/// `launch.model` is the model's manifest, already on the disk, and `chosen` what libwaifu says +/// of it. `why` is what `auto` said about the device it picked, if it picked one; `machine` is +/// what the device row offers instead. `image` fills the draw from box, where `-i` named one, and +/// `output` is the directory every picture is written into. pub fn draw( launch: &Launch, why: Option, - manifest: &Path, + chosen: &Chosen, image: Option<&str>, output: &Path, machine: &Machine, ) -> Result<(), Error> { - // Read here, while nothing is being waited for: the manifest is a few kilobytes, and the model - // behind it is minutes of reading. What is wrong with it is reported once, by the painter, - // with the error that says why. - let chosen = describe::look_at(&launch.model); - let label = label_of(&launch.model); + let manifest = PathBuf::from(&launch.model); + let label = launch.label.clone(); // Set while a run is in flight to ask it to stop between steps, which is the only place it // can be asked: a step, once started, is a kernel launch that nothing here can call back. @@ -109,16 +103,15 @@ pub fn draw( let painter = std::thread::spawn({ let cancel = Arc::clone(&cancel); - let manifest = manifest.to_path_buf(); let output = output.to_path_buf(); let label = label.clone(); move || paint(&manifest, &label, &output, &waiting, &updates, &cancel) }); - let mut app = App::for_chosen(launch.task, label, launch.runtime, &chosen, output); + let mut app = App::for_chosen(launch.task, label, launch.runtime, chosen, output); app.available = machine.available.clone(); app.room = machine.room; - app.weights = hub::model_bytes(&launch.model); + app.weights = chosen.bytes; // -i only fills the box. Everything about a run is changeable between runs, and a picture // named on the command line is no different -- it is where to start, not what to be stuck @@ -137,10 +130,10 @@ pub fn draw( app.after_reading = (!notes.is_empty()).then(|| notes.join("; ")); app.read_onto(launch.runtime); - // The tensor library's lines go to the screen rather than over it from here on: what it found - // on a device it is asked about for the first time is written to the console otherwise, and the - // console is under the screen. A warning or an error is put on the status line; the rest is - // let go of, since nothing on this screen would read it. + // The libraries' lines go to the screen rather than over it from here on: what the tensor + // library found on a device it is asked about for the first time is written to the console + // otherwise, and the console is under the screen. A warning or an error is put on the status + // line; the rest is let go of, since nothing on this screen would read it. let heard: Heard = Arc::default(); crate::log::set_sink(Some(Box::new({ let heard = Arc::clone(&heard); @@ -165,17 +158,18 @@ pub fn draw( ratatui::restore(); crate::log::set_sink(None); - // Waited for rather than left where it is. Under `Residency::LowVram` the model is some - // hundreds of page-locked host allocations, and page-locked memory is the CUDA driver's rather + // Waited for rather than left where it is. Offloaded to the host, the model is some hundreds + // of page-locked host allocations, and page-locked memory is the CUDA driver's rather // than the operating system's. A process that exits with this thread still holding it has two // teardowns running at once -- the driver unloading, and this thread dropping the model as its // channel closes -- and the one that loses gives its memory back to a driver that is already // gone. // - // So the order is said instead of raced: the run is asked to stop, the channel it reads is - // closed, and the model is let go of before this function returns. The wait is bounded by the - // step the run is on -- `cancel` is already set, so it stops at the next one -- except for a - // decode, or a model still being read, neither of which can be cut short. + // So the order is said instead of raced: the run is asked to stop, the channel the painter + // reads is closed, and the engine -- and the model with it -- is freed before this function + // returns. The wait is bounded by the step the run is on -- `cancel` is already set, so it + // stops at the next one -- except for a decode, or a model still being read, neither of which + // can be cut short. cancel.store(true, Ordering::Relaxed); drop(requests); let _ = painter.join(); @@ -186,16 +180,6 @@ pub fn draw( /// What the tensor library has said that is worth putting on the screen, waiting to be. type Heard = Arc>>; -/// What the screens call a model: as typed when it was a name, because that is what someone -/// would type again, and the file name when it was a path, because the whole path does not fit -/// and its tail is the part that identifies it. -fn label_of(model: &str) -> String { - match model.rsplit_once(['/', '\\']) { - Some((_, file)) if !file.is_empty() => file.to_string(), - _ => model.to_string(), - } -} - /// Draws, reads the keyboard, and passes what the two channels carry between them. fn run( mut terminal: DefaultTerminal, @@ -315,8 +299,8 @@ enum Update { Failed(String), } -/// The thread that owns the model: it reads it onto a device when asked, runs jobs on it, and -/// says what it is doing as it does it. +/// The thread that talks to the engine: it starts one on a device and has it read the model when +/// asked, runs jobs on it, and says what it is doing as it does it. fn paint( manifest: &Path, label: &str, @@ -325,24 +309,24 @@ fn paint( updates: &Sender, cancel: &AtomicBool, ) { - let mut model: Option = None; + let mut engine: Option = None; for request in requests { let job = match request { Request::Read(runtime) => { // Let go of first. Two copies on one card is twice the memory, and a move that // ran out of it would lose the model that was working as well as the new one. - drop(model.take()); - let update = match ImageModel::from_manifest( - manifest, - runtime.device(), - runtime.residency(), - ) { - Ok(read) => { - model = Some(read); - Update::Ready - } - Err(error) => Update::Failed(error.to_string()), + drop(engine.take()); + let update = match Engine::new(runtime) { + Ok(started) => match started.load(manifest) { + Ended::Done(()) => { + engine = Some(started); + Update::Ready + } + Ended::Cancelled => Update::Failed("the read was stopped".to_string()), + Ended::Failed(error) => Update::Failed(error), + }, + Err(error) => Update::Failed(error), }; if updates.send(update).is_err() { return; @@ -353,7 +337,7 @@ fn paint( }; // The screen does not send one while there is no model, but a job is a job. - let Some(model) = &model else { + let Some(engine) = &engine else { let _ = updates.send(Update::Failed("there is no model to draw with".to_string())); continue; }; @@ -362,15 +346,6 @@ fn paint( cancel.store(false, Ordering::Relaxed); let started = Instant::now(); - let mut report = |progress| { - let _ = updates.send(Update::Progress(progress)); - if cancel.load(Ordering::Relaxed) { - ControlFlow::Break(()) - } else { - ControlFlow::Continue(()) - } - }; - // A picture to start from is read and scaled here, before the run, so that a path that // is not there is one message rather than a run that fails partway. let started_from = match &job.from { @@ -384,30 +359,34 @@ fn paint( None => None, }; - let drawn = match &started_from { - Some(image) => { - model.generate_from_image_reporting(image, &job.prompt, &job.options, &mut report) - } - None => model.generate_reporting(&job.prompt, &job.options, &mut report), + let (say, news) = channel::(); + let from = started_from + .as_ref() + .map(|(width, height, rgb)| (*width, *height, rgb.as_slice())); + let drawn = match engine.draw(&job.prompt, &job.options, from, &say) { + Ok(id) => wait_for_the_picture(engine, id, &news, updates, cancel), + Err(error) => Ended::Failed(error), }; let update = match drawn { // Written here rather than back on the drawing thread: the pixels are three megabytes // that nothing on the other side would do anything with but hand to a file. - Ok(Some(image)) => match keep(&image, output, &job.parameters(label)) { - Ok((path, width, height)) => Update::Done { - kept: Kept { - path, - width, - height, - seed: job.options.seed.unwrap_or_default(), + Ended::Done((width, height, rgb)) => { + match keep(width, height, &rgb, output, &job.parameters(label)) { + Ok(path) => Update::Done { + kept: Kept { + path, + width: width as usize, + height: height as usize, + seed: job.options.seed.unwrap_or_default(), + }, + elapsed: started.elapsed(), }, - elapsed: started.elapsed(), - }, - Err(error) => Update::Failed(error.to_string()), - }, - Ok(None) => Update::Stopped, - Err(error) => Update::Failed(error.to_string()), + Err(error) => Update::Failed(error.to_string()), + } + } + Ended::Cancelled => Update::Stopped, + Ended::Failed(error) => Update::Failed(error), }; if updates.send(update).is_err() { @@ -416,7 +395,37 @@ fn paint( } } -/// A picture from a file, as the `(1, 3, height, width)` tensor a run can start from. +/// Passes a run's progress on to the screen until it ends, stopping it once when `cancel` is set, +/// and hands back how it ended. `say`, the sending end of `news`, lives in the caller until then: +/// the engine's callbacks send down it, and exactly one end comes for a job that was started. +fn wait_for_the_picture( + engine: &Engine, + job: u64, + news: &Receiver, + updates: &Sender, + cancel: &AtomicBool, +) -> Ended<(u32, u32, Vec)> { + let mut asked_to_stop = false; + loop { + match news.recv_timeout(Duration::from_millis(100)) { + Ok(News::Progress(progress)) => { + let _ = updates.send(Update::Progress(progress)); + } + Ok(News::Drawn(ended)) => return ended, + Ok(News::Loaded(_)) => {} + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {} + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => { + return Ended::Failed("the engine went away".to_string()); + } + } + if !asked_to_stop && cancel.load(Ordering::Relaxed) { + engine.cancel(job); + asked_to_stop = true; + } + } +} + +/// A picture from a file, as the rows of RGB a run can start from, and its size. /// /// Scaled to the size the screen asks for rather than kept at its own: the model works at /// multiples of 32, and a picture off a camera or a phone is not one. Stretched to it rather than @@ -426,7 +435,7 @@ fn paint( /// Lanczos, because this is the direction that loses pixels -- a photograph is larger than /// anything SDXL draws -- and a cheaper filter leaves stair steps that the run then faithfully /// keeps. -fn read_image(path: &Path, width: i32, height: i32) -> Result { +fn read_image(path: &Path, width: i32, height: i32) -> Result<(u32, u32, Vec), Error> { let opened = image::ImageReader::open(path)? .with_guessed_format()? .decode()?; @@ -436,19 +445,20 @@ fn read_image(path: &Path, width: i32, height: i32) -> Result { image::imageops::FilterType::Lanczos3, ); - Ok(from_rgb8(width, height, scaled.to_rgb8().as_raw())?) + Ok((width as u32, height as u32, scaled.to_rgb8().into_raw())) } -/// Writes the tensor a run ends with into the first `waifu-NNNN.png` in `directory` that nothing -/// else has, with `parameters` written into it, and says where and how big. +/// Writes the picture a run ends with -- rows of RGB, `width` by `height` -- into the first +/// `waifu-NNNN.png` in `directory` that nothing else has, with `parameters` written into it, and +/// says where. fn keep( - image: &Tensor, + width: u32, + height: u32, + rgb: &[u8], directory: &Path, parameters: &str, -) -> Result<(PathBuf, usize, usize), Error> { - let shape = image.shape(); - let (width, height) = (shape[3] as usize, shape[2] as usize); - let bytes = png(image, parameters)?; +) -> Result { + let bytes = png(width, height, rgb, parameters)?; for number in 1..10_000 { let path = directory.join(format!("waifu-{number:04}.png")); @@ -466,7 +476,7 @@ fn keep( }; file.write_all(&bytes)?; - return Ok((path, width, height)); + return Ok(path); } Err(format!( @@ -476,17 +486,13 @@ fn keep( .into()) } -/// The tensor a run ends with, as a PNG, with `parameters` written into it. -fn png(image: &Tensor, parameters: &str) -> Result, Error> { +/// The picture a run ends with, as a PNG, with `parameters` written into it. +fn png(width: u32, height: u32, pixels: &[u8], parameters: &str) -> Result, Error> { use image::ImageEncoder; - let pixels = to_rgb8(image)?; - let shape = image.shape(); - let (width, height) = (shape[3] as u32, shape[2] as u32); - let mut encoded = Vec::new(); image::codecs::png::PngEncoder::new(&mut encoded).write_image( - &pixels, + pixels, width, height, image::ExtendedColorType::Rgb8, @@ -867,7 +873,7 @@ impl App { task, model, runtime, - output: PathBuf::from(crate::cli::args::DEFAULT_OUTPUT), + output: PathBuf::from(crate::args::DEFAULT_OUTPUT), prompt: TextField::default(), negative: TextField::default(), from: TextField::default(), @@ -948,7 +954,7 @@ impl App { )); } - /// The screen for a model as [`describe::look_at`] describes it: its numbers in the boxes, its + /// The screen for a model as libwaifu describes it: its numbers in the boxes, its /// sizes in the list, and what its card suggests typing already typed. fn for_chosen( task: Task, @@ -1995,7 +2001,7 @@ mod tests { use ratatui::backend::TestBackend; use super::*; - use crate::describe::SIZES; + use crate::engine::SIZES; fn press(code: KeyCode) -> KeyEvent { KeyEvent::new(code, KeyModifiers::NONE) @@ -2093,10 +2099,10 @@ mod tests { /// The screen for the model a manifest describes, as `draw` opens it. fn for_manifest(path: &Path) -> App { - let chosen = describe::look_at(path.to_str().unwrap()); + let chosen = crate::engine::look_at(path).unwrap(); App::for_chosen( Task::Txt2Img, - label_of(path.to_str().unwrap()), + crate::draw::label_of(path.to_str().unwrap()), Runtime::ALL[0], &chosen, Path::new("out"), @@ -2140,6 +2146,7 @@ mod tests { #[test] fn a_path_is_called_by_its_file_name_and_a_name_by_itself() { + use crate::draw::label_of; assert_eq!(label_of("sdxl:base"), "sdxl:base"); assert_eq!(label_of("/models/mine/sdxl.yaml"), "sdxl.yaml"); assert_eq!(label_of("C:\\models\\sdxl.yaml"), "sdxl.yaml"); @@ -2857,10 +2864,9 @@ mod tests { } #[test] - fn a_picture_off_the_disk_becomes_a_tensor_at_the_size_asked_for() { + fn a_picture_off_the_disk_becomes_rgb_at_the_size_asked_for() { // Through a real file rather than a buffer: the point of this one is that what the image - // crate decodes, scales and hands back lines up with what from_rgb8 expects, which a - // tensor built by hand would not check. + // crate decodes and scales comes out as the rows of RGB the engine reads. let path = std::env::temp_dir().join(format!("waifu-read-image-{}.png", std::process::id())); let pixels: Vec = (0..32 * 32) @@ -2870,24 +2876,16 @@ mod tests { // A size that is neither the file's nor a whole multiple of it, since the run decides the // size and a picture off a camera never happens to be it. - let image = read_image(&path, 192, 128).unwrap(); + let (width, height, rgb) = read_image(&path, 192, 128).unwrap(); std::fs::remove_file(&path).unwrap(); - assert_eq!(image.shape(), vec![1, 3, 128, 192]); - - // Scaling is allowed to move a value; leaving the range the model reads is not. - let values = image.to_vec_f32().unwrap(); - assert!( - values.iter().all(|x| (-1.0..=1.0).contains(x)), - "a scaled picture left [-1, 1]" - ); + assert_eq!((width, height), (192, 128)); + assert_eq!(rgb.len(), 192 * 128 * 3); // The third channel was one value everywhere, so scaling cannot have made it anything - // else -- which is the check that the planes did not get shuffled on the way through. - let plane = 128 * 192; - let blue = 7.0 / 127.5 - 1.0; + // else -- which is the check that the channels did not get shuffled on the way through. assert!( - values[2 * plane..].iter().all(|x| (x - blue).abs() < 1e-3), + rgb.chunks(3).all(|pixel| pixel[2].abs_diff(7) <= 1), "the channels came back in the wrong order" ); } @@ -3205,11 +3203,10 @@ mod tests { // One already there, which is not written over. std::fs::write(directory.join("waifu-0001.png"), b"somebody else's").unwrap(); - let image = from_rgb8(2, 2, &[255, 0, 0, 0, 255, 0, 0, 0, 255, 9, 9, 9]).unwrap(); - let (path, width, height) = keep(&image, &directory, "a cat\nSeed: 7").unwrap(); + let image = [255, 0, 0, 0, 255, 0, 0, 0, 255, 9, 9, 9]; + let path = keep(2, 2, &image, &directory, "a cat\nSeed: 7").unwrap(); assert_eq!(path, directory.join("waifu-0002.png")); - assert_eq!((width, height), (2, 2)); assert_eq!( std::fs::read(directory.join("waifu-0001.png")).unwrap(), b"somebody else's" @@ -3316,7 +3313,7 @@ mod tests { // Both halves, because either alone is useless: the name without the revision does not // say which code, and a bare hash on a screenshot does not say what it is a hash of. assert!(drawn.contains("libwaifu"), "{drawn}"); - assert!(drawn.contains(crate::cli::REVISION), "{drawn}"); + assert!(drawn.contains(crate::REVISION), "{drawn}"); } #[test] diff --git a/waifu/src/cli/tui/field.rs b/waifu-cli/src/tui/field.rs similarity index 100% rename from waifu/src/cli/tui/field.rs rename to waifu-cli/src/tui/field.rs diff --git a/waifu/src/cli/tui/files.rs b/waifu-cli/src/tui/files.rs similarity index 99% rename from waifu/src/cli/tui/files.rs rename to waifu-cli/src/tui/files.rs index 59a23f3..4af73d2 100644 --- a/waifu/src/cli/tui/files.rs +++ b/waifu-cli/src/tui/files.rs @@ -40,7 +40,7 @@ use ratatui::text::{Line, Span}; use ratatui::widgets::{Block, BorderType, Clear, Paragraph}; use ratatui::Frame; -use crate::cli::tui::centred; +use crate::tui::centred; /// The name a parent directory goes by, and the row that walks to it. const PARENT: &str = ".."; diff --git a/waifu/src/cli/tui/mod.rs b/waifu-cli/src/tui/mod.rs similarity index 91% rename from waifu/src/cli/tui/mod.rs rename to waifu-cli/src/tui/mod.rs index 5c96523..8c5ad72 100644 --- a/waifu/src/cli/tui/mod.rs +++ b/waifu-cli/src/tui/mod.rs @@ -47,9 +47,8 @@ use ratatui::text::{Line, Span}; use ratatui::widgets::{Block, BorderType}; use ratatui::{DefaultTerminal, Frame}; -use crate::cli::args::{DeviceOption, Room, Runtime}; -use crate::cli::hub; -use crate::cli::task::Task; +use crate::args::{DeviceOption, Room, Runtime}; +use crate::task::Task; type Error = Box; @@ -80,26 +79,17 @@ pub fn can_take_the_screen() -> bool { io::stdin().is_terminal() && io::stdout().is_terminal() } -/// Has the terminal put back the way it was if the tensor library gives up, from now on. -/// -/// Called before the first screen goes up and left for the rest of the run: every screen takes -/// the terminal, and a check that fails inside the tensor library can come from any of them -- -/// the drawing screen has a model being read or run behind it the whole time it is up. -pub fn give_the_screen_back_if_flint_aborts() { - crate::flint::on_fatal(give_the_screen_back); -} - /// Puts the terminal back the way it was, from inside the tensor library's last breath. /// -/// A check that fails there prints what went wrong and calls `abort()`, and the message is worth +/// A check that fails there writes a fatal line and calls `abort()`, and the message is worth /// nothing written across a screen full of boxes -- which is where it lands, since by then the /// alternate screen is up, the cursor is hidden and the terminal is in raw mode. So the last thing -/// that happens before the message is this. +/// that happens before the line is written is this: see [`crate::log`]. /// /// It runs on whichever thread failed, which may not be the one drawing. Racing a half-finished /// frame is not a concern: nothing after this draws another, and a frame torn on the way out is /// better than a message nobody can read. -extern "C" fn give_the_screen_back() { +pub fn give_the_screen_back() { ratatui::restore(); // And the cursor, which restore() does not do: hiding it is the Terminal's and showing it @@ -138,8 +128,7 @@ impl Machine { /// second-guessed, but a run on it that is going to run out of memory is still worth a warning /// before it is started. pub fn measure(device: DeviceOption) -> Machine { - let available = Runtime::available(); - let room = Room::measure(&available); + let (available, room) = crate::runtime::measure(); Machine { device, available, @@ -147,14 +136,13 @@ impl Machine { } } - /// Where `model` runs to begin with, and -- where that was worked out from the size of the - /// model -- the line that says so. - pub fn start_for(&self, model: &str) -> (Runtime, Option) { + /// Where a model of `weights` bytes runs to begin with, and -- where that was worked out from + /// the size of the model -- the line that says so. + pub fn start_for(&self, weights: Option) -> (Runtime, Option) { if self.device != DeviceOption::Auto { - return (self.device.resolve(), None); + return (self.device.resolve(&self.available), None); } - let weights = hub::model_bytes(model); let runtime = Runtime::best(weights, &self.available, self.room); let why = weights.map(|weights| { // CUDA's room where there is CUDA: that is the card both of its rows are about. @@ -224,7 +212,7 @@ fn heading(trail: &[&str]) -> Line<'static> { let mut spans = vec![ Span::raw(" libwaifu").bold(), Span::raw(" "), - Span::styled(crate::cli::REVISION, Style::new().fg(Color::DarkGray)), + Span::styled(crate::REVISION, Style::new().fg(Color::DarkGray)), ]; for step in trail { spans.push(Span::styled(" > ", Style::new().fg(Color::DarkGray))); @@ -360,7 +348,7 @@ mod tests { let drawn = screen_text(&terminal); assert!(drawn.contains("libwaifu"), "{drawn}"); - assert!(drawn.contains(crate::cli::REVISION), "{drawn}"); + assert!(drawn.contains(crate::REVISION), "{drawn}"); assert!(drawn.contains("txt2img > sdxl:base"), "{drawn}"); } diff --git a/waifu/src/cli/tui/models.rs b/waifu-cli/src/tui/models.rs similarity index 92% rename from waifu/src/cli/tui/models.rs rename to waifu-cli/src/tui/models.rs index dcbcfa5..5bab291 100644 --- a/waifu/src/cli/tui/models.rs +++ b/waifu-cli/src/tui/models.rs @@ -46,12 +46,11 @@ use ratatui::text::{Line, Span}; use ratatui::widgets::Paragraph; use ratatui::{DefaultTerminal, Frame}; -use crate::cli::config; -use crate::cli::hub::{self, Said}; -use crate::cli::task::Task; -use crate::cli::tui::ask::{Answer, Confirm, Input}; -use crate::cli::tui::files::{self, tail, FilePicker}; -use crate::cli::tui::{bar, bordered, gigabytes, heading, quits, row_style, Picked, Step}; +use crate::manager::{self, Said}; +use crate::task::Task; +use crate::tui::ask::{Answer, Confirm, Input}; +use crate::tui::files::{self, tail, FilePicker}; +use crate::tui::{bar, bordered, gigabytes, heading, quits, row_style, Picked, Step}; type Error = Box; @@ -128,18 +127,18 @@ impl Entry { /// The published models a task can run, in the catalogue's order, and the file row under them. fn entries_for(task: Task) -> Vec { - let mut entries: Vec = hub::listed() + let mut entries: Vec = manager::pictures() .into_iter() // What img2img would only refuse. Asked of the kind and of the manifest where it is here, // which is as much as can be known without reading the weights. - .filter(|model| task != Task::Img2Img || crate::cli::draw::draws_from_a_picture(model.name)) + .filter(|model| task != Task::Img2Img || crate::draw::draws_from_a_picture(model.name)) .map(|model| Entry::Published { name: model.name, full_name: model.full_name, cached: model.cached, bytes: model.bytes, explicit: model.explicit, - at: hub::local_directory(model.name), + at: model.directory, }) .collect(); @@ -248,6 +247,7 @@ impl Choices { /// every row says. fn read_the_cache_again(&mut self) { self.directory = model_directory(); + let catalog = manager::catalog(); for entry in &mut self.entries { if let Entry::Published { name, @@ -257,27 +257,33 @@ impl Choices { .. } = entry { - *cached = hub::is_cached(name); - *bytes = hub::cached_bytes(name); - *at = hub::local_directory(name); + let listed = catalog.iter().find(|listed| listed.name == *name); + *cached = listed.is_some_and(|listed| listed.cached); + *bytes = listed.map_or(0, |listed| listed.bytes); + *at = listed.and_then(|listed| listed.directory.clone()); } } } } -/// [`hub::model_directory`], with its error said the way the screen says one. +/// Where models are fetched into and who said so, as the model manager says. fn model_directory() -> Result<(PathBuf, Said), String> { - hub::model_directory().map_err(|error| error.to_string()) + manager::directory().map(|directory| (directory.now, directory.said)) } /// The box that changes where models are fetched into, holding the directory as it stands. fn ask_for_a_directory(now: &Path) -> Input { - let file = config::path() + let directory = manager::directory().ok(); + let file = directory + .as_ref() + .and_then(|directory| directory.config_file.as_ref()) .map(|file| file.display().to_string()) - .unwrap_or_else(|| config::FILE_NAME.to_string()); - let default = hub::default_directory() - .map(|directory| directory.display().to_string()) - .unwrap_or_else(|_| "the default".to_string()); + .unwrap_or_else(|| manager::CONFIG_FILE_NAME.to_string()); + let default = directory + .as_ref() + .and_then(|directory| directory.default.as_ref()) + .map(|default| default.display().to_string()) + .unwrap_or_else(|| "the default".to_string()); Input::ask("where models are fetched into", &now.display().to_string()).saying(&format!( "Models fetched from now on go here; ones already fetched stay where they are. Kept in \ @@ -304,17 +310,13 @@ fn typed_directory(typed: &str) -> Result { /// first download into it. fn change_directory(typed: &str) -> Result<(), String> { if typed.is_empty() { - return config::set_model_dir(None) - .map(|_| ()) - .map_err(|error| error.to_string()); + return manager::set_directory(None); } let directory = typed_directory(typed)?; std::fs::create_dir_all(&directory) .map_err(|error| format!("cannot make {}: {error}", directory.display()))?; - config::set_model_dir(Some(&directory)) - .map(|_| ()) - .map_err(|error| error.to_string()) + manager::set_directory(Some(&directory)) } /// Runs the list until a model is chosen and on the disk, or somebody goes back or leaves. @@ -407,7 +409,7 @@ pub fn choose(terminal: &mut DefaultTerminal, task: Task) -> Result } (Answer::Given(true), Asking::Delete(name)) => { asking = None; - match hub::remove(name) { + match manager::remove(name) { // Read back off the disk rather than assumed: what the row says about a // model is what is there, and this is the same call that first said it. Ok(()) => choices.read_the_cache_again(), @@ -518,9 +520,11 @@ pub fn choose(terminal: &mut DefaultTerminal, task: Task) -> Result Ok((now, _)) => editing = Some(ask_for_a_directory(now)), // Asked anyway: a file that does not read is mended by writing a good one. Err(_) => { - editing = Some(ask_for_a_directory( - &hub::default_directory().unwrap_or_default(), - )) + let default = manager::directory() + .ok() + .and_then(|directory| directory.default) + .unwrap_or_default(); + editing = Some(ask_for_a_directory(&default)) } } } @@ -598,46 +602,57 @@ fn start_fetching(name: &'static str) -> Doing { let (say, news) = channel::(); let stop = Arc::new(AtomicBool::new(false)); + // A thread between the screen and the fetch, which passes its news on as the screen reads it + // and asks it to stop when the screen says so. The fetch is the model manager's, on a thread + // of its own; this one only listens to it. std::thread::spawn({ let stop = Arc::clone(&stop); move || { - let mut report = |progress: hub::Progress| { - let word = match progress { - hub::Progress::From { hub } => Word::From(hub), - hub::Progress::Fetching { + let fetch = match manager::Fetch::start(name) { + Ok(fetch) => fetch, + Err(error) => { + let _ = say.send(Word::Failed(error)); + return; + } + }; + let mut asked_to_stop = false; + let mut hub_said = false; + let word = loop { + match fetch.news.recv_timeout(TICK) { + Ok(manager::Word::Fetching { + hub, file, done, total, part, parts, - } => Word::Fetching { - file: file.to_string(), - done, - total, - part, - parts, - }, - hub::Progress::Fetched { - file, - bytes, - part, - parts, - } => Word::Fetching { - file: file.to_string(), - done: bytes, - total: Some(bytes), - part, - parts, - }, - }; - let _ = say.send(word); - }; - - let stopping = || stop.load(Ordering::Relaxed); - let word = match hub::resolve_reporting(name, &mut report, &stopping) { - Ok(_) => Word::Done, - Err(error) if hub::stopped(&error) => Word::Stopped(error.to_string()), - Err(error) => Word::Failed(error.to_string()), + }) => { + if let (false, Some(hub)) = (hub_said, hub) { + hub_said = true; + let _ = say.send(Word::From(manager::kept(&hub))); + } + if !file.is_empty() { + let _ = say.send(Word::Fetching { + file, + done, + total, + part, + parts, + }); + } + } + Ok(manager::Word::Done) => break Word::Done, + Ok(manager::Word::Stopped(said)) => break Word::Stopped(said), + Ok(manager::Word::Failed(said)) => break Word::Failed(said), + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {} + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => { + break Word::Failed("the fetch stopped without saying why".to_string()) + } + } + if !asked_to_stop && stop.load(Ordering::Relaxed) { + fetch.stop(); + asked_to_stop = true; + } }; let _ = say.send(word); } @@ -974,7 +989,7 @@ mod tests { use ratatui::backend::TestBackend; use super::*; - use crate::cli::tui::screen_text; + use crate::tui::screen_text; /// The foot, drawn into a buffer, as one long string of what it says. fn foot(doing: &Doing) -> String { diff --git a/waifu/src/cli/tui/tasks.rs b/waifu-cli/src/tui/tasks.rs similarity index 98% rename from waifu/src/cli/tui/tasks.rs rename to waifu-cli/src/tui/tasks.rs index 0e5da8b..34b81e1 100644 --- a/waifu/src/cli/tui/tasks.rs +++ b/waifu-cli/src/tui/tasks.rs @@ -30,9 +30,9 @@ use ratatui::text::{Line, Span}; use ratatui::widgets::Paragraph; use ratatui::{DefaultTerminal, Frame}; -use crate::cli::task::Task; -use crate::cli::tui::ask::Answer; -use crate::cli::tui::{bordered, heading, quits, Step}; +use crate::task::Task; +use crate::tui::ask::Answer; +use crate::tui::{bordered, heading, quits, Step}; type Error = Box; @@ -147,7 +147,7 @@ mod tests { use ratatui::backend::TestBackend; use ratatui::crossterm::event::KeyModifiers; - use crate::cli::tui::screen_text; + use crate::tui::screen_text; fn press(code: KeyCode) -> KeyEvent { KeyEvent::new(code, KeyModifiers::NONE) diff --git a/waifu-ffi/Cargo.toml b/waifu-ffi/Cargo.toml deleted file mode 100644 index ceac459..0000000 --- a/waifu-ffi/Cargo.toml +++ /dev/null @@ -1,19 +0,0 @@ -[package] -name = "waifu-ffi" -version = "0.1.0" -edition = "2021" -license = "MIT" -description = "The C API over libwaifu: what Swift imports and .NET calls through P/Invoke" - -# A dynamic library, libwaifu.dylib once CMake has named it: the Metal kernels are a section of -# the image flint is in, and a dylib carries them wherever it goes. An rlib as well, for the tests, -# which call the C functions the way C would. -[lib] -name = "waifu_ffi" -crate-type = ["cdylib", "rlib"] - -[dependencies] -waifu = { path = "../waifu", features = ["engine"] } - -[build-dependencies] -cbindgen = { version = "0.29", default-features = false } diff --git a/waifu-modelmanager/Cargo.toml b/waifu-modelmanager/Cargo.toml new file mode 100644 index 0000000..d2438f3 --- /dev/null +++ b/waifu-modelmanager/Cargo.toml @@ -0,0 +1,59 @@ +[package] +name = "waifu-modelmanager" +version = "0.1.0" +edition = "2021" +license = "MIT" +description = "libwaifu_modelmanager: which models there are, fetching them, and where they are kept" + +# A dynamic library, libwaifu_modelmanager.dylib (.so, waifu_modelmanager.dll), whose C API is all +# anything outside this crate uses: see docs/ffi.md. The rlib is for this crate's own tests, for +# the command line, which links it statically and calls its C API, and for waifu-tools. libwaifu does not link it, nor it libwaifu. +[lib] +name = "waifu_modelmanager" +path = "src/lib.rs" +crate-type = ["cdylib", "rlib"] + +[dependencies] + +# The catalogue as JSON, for the C API to hand on: a parser this crate would own for no reason, +# since it is already in the tree, under hf-hub. +# +# `Value` rather than derived types. What crosses is a handful of shapes, each built once, and a +# derive would buy nothing at the cost of a second description of every one of them. +serde_json = { version = "1" } + +# The settings file, `config.toml`, which today says one thing: where fetched models go. It is a +# file for people to open in an editor, so it is in the format that lets them write a comment and +# a Windows path without escaping every backslash -- and reading that properly is a parser, which +# is not something to write again here. +toml = { version = "1" } + +# Fetching from ModelScope, which hf-hub does not speak, and probing whether the machine can +# reach the outside at all. Rust has no HTTP in its standard library, so this is the one thing the +# tool cannot do on its own. Held to the rustls backend: the native-tls default would link the +# system OpenSSL, and a dynamic library on the far side of the build is exactly what this crate +# does not want. +ureq = { version = "3", default-features = false, features = [ + "rustls", + "gzip", +] } + +# Fetching from Hugging Face, which is where the models are published and where all but one of +# the fetches go. Its own client rather than ureq against a URL built here: it is what tracks how +# that hub actually answers -- redirects to a CDN, xet-backed files that are not a plain GET, the +# headers that say a file moved -- and none of that is worth rediscovering here one bug report at +# a time. +# +# It is the largest thing in this file by a distance, reqwest and tokio and hf-xet behind it, and +# that is the price of the above. Held to rustls for the same reason ureq is, and to the blocking +# feature because nothing else in this crate is async: an executor exists to run one download at +# a time and no further. +# +# The other fetch is ModelScope, which this does not speak. That one stays on ureq, above. +hf-hub = { version = "1", default-features = false, features = [ + "blocking", + "rustls-tls", +] } + +[build-dependencies] +cbindgen = { version = "0.29", default-features = false } diff --git a/waifu-modelmanager/build.rs b/waifu-modelmanager/build.rs new file mode 100644 index 0000000..9315e87 --- /dev/null +++ b/waifu-modelmanager/build.rs @@ -0,0 +1,62 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! Names the library for where it is loaded from, and writes waifu_modelmanager.h. +//! +//! Nothing native is linked: the model manager is Rust all the way down, so this builds without +//! CMake having built anything first. + +use std::path::{Path, PathBuf}; + +fn main() { + let manifest = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + + // Found next to whatever loads it -- an app's Frameworks directory, a test's own -- rather than + // at the path cargo happened to write it to. + if std::env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("macos") { + println!("cargo:rustc-cdylib-link-arg=-Wl,-install_name,@rpath/libwaifu_modelmanager.dylib"); + } + + write_header(&manifest); +} + +/// waifu_modelmanager.h, from the `extern "C"` functions and `#[repr(C)]` types in src/c_api.rs and +/// src/c_api/support.rs, with their doc comments. Written beside the source, where a change to the +/// ABI shows up in a diff. +fn write_header(manifest: &Path) { + println!("cargo:rerun-if-changed=src/c_api.rs"); + println!("cargo:rerun-if-changed=src/c_api/support.rs"); + println!("cargo:rerun-if-changed=cbindgen.toml"); + + let config = cbindgen::Config::from_file(manifest.join("cbindgen.toml")) + .expect("waifu-modelmanager/cbindgen.toml"); + match cbindgen::Builder::new() + .with_config(config) + .with_src(manifest.join("src/c_api.rs")) + .with_src(manifest.join("src/c_api/support.rs")) + .generate() + { + Ok(bindings) => { + bindings.write_to_file(manifest.join("include/waifu_modelmanager.h")); + } + // Said rather than failed on: a header that cannot be regenerated leaves the last one, + // and the library under it still builds. + Err(error) => println!("cargo:warning=waifu_modelmanager.h was not regenerated: {error}"), + } +} diff --git a/waifu-modelmanager/cbindgen.toml b/waifu-modelmanager/cbindgen.toml new file mode 100644 index 0000000..88b8d14 --- /dev/null +++ b/waifu-modelmanager/cbindgen.toml @@ -0,0 +1,33 @@ +# How waifu_modelmanager.h is written from src/c_api.rs. See docs/ffi.md for the contract it keeps. +# +# Everything the header declares is this library's own: it includes nothing but the C library's, +# and shares no type with libwaifu's. + +language = "C" +include_guard = "LIBWAIFU_WAIFU_MODELMANAGER_H" +cpp_compat = true +documentation = true +documentation_style = "c99" +style = "type" +sys_includes = ["stdbool.h", "stddef.h", "stdint.h"] +no_includes = true +# size_t for a length, as C callers expect and Swift imports as Int. +usize_is_size_t = true +autogen_warning = "/* Written by cbindgen from waifu-modelmanager/src/c_api.rs. Change that, not this. */" +header = """/* The MIT License (MIT). Copyright (c) 2026 Xiaoyang Chen. See LICENSE. */""" + +[defines] + +[enum] +rename_variants = "None" +prefix_with_name = false + +[export] +include = [ + "WaifuModelManagerStatusCode", + "WaifuModelManagerStatus", + "WaifuModelManagerLogLevel", + "WaifuModelManagerLogEvent", + "WaifuModelManagerProgressEvent", + "WaifuModelManagerCompletionEvent", +] diff --git a/waifu-modelmanager/include/waifu_modelmanager.h b/waifu-modelmanager/include/waifu_modelmanager.h new file mode 100644 index 0000000..b4af6d3 --- /dev/null +++ b/waifu-modelmanager/include/waifu_modelmanager.h @@ -0,0 +1,180 @@ +/* The MIT License (MIT). Copyright (c) 2026 Xiaoyang Chen. See LICENSE. */ + +#ifndef LIBWAIFU_WAIFU_MODELMANAGER_H +#define LIBWAIFU_WAIFU_MODELMANAGER_H + +/* Written by cbindgen from waifu-modelmanager/src/c_api.rs. Change that, not this. */ + +#include +#include +#include + +// What a caller compiled against this header expects. Bumped when anything in it changes in a +// way an older caller would misread. +#define WAIFU_MODELMANAGER_ABI_VERSION 1 + +typedef enum { + WAIFU_MODELMANAGER_LOG_DEBUG = 0, + WAIFU_MODELMANAGER_LOG_INFO = 1, + WAIFU_MODELMANAGER_LOG_WARNING = 2, + WAIFU_MODELMANAGER_LOG_ERROR = 3, + // The last line: the process aborts when the callback returns. + WAIFU_MODELMANAGER_LOG_FATAL = 4, +} WaifuModelManagerLogLevel; + +// What a call or a job came to. +typedef enum { + WAIFU_MODELMANAGER_OK = 0, + WAIFU_MODELMANAGER_CANCELLED = 1, + // A NULL, or a string that is not UTF-8. + WAIFU_MODELMANAGER_ERR_INVALID_ARGUMENT = 2, + // No published model by that name, or no manifest at that path. + WAIFU_MODELMANAGER_ERR_UNKNOWN_MODEL = 3, + // The download failed. + WAIFU_MODELMANAGER_ERR_FETCH = 4, + // A panic in the library, or a file that could not be read or written. + WAIFU_MODELMANAGER_ERR_INTERNAL = 5, +} WaifuModelManagerStatusCode; + +// A download running on a thread of its own. +typedef struct WaifuModelFetch WaifuModelFetch; + +typedef struct { + WaifuModelManagerLogLevel level; + // Where it was written: a module's name. + const char *source; + // The line alone, with no level or time in front of it. + const char *message; +} WaifuModelManagerLogEvent; + +// Called on whichever thread wrote the line, which can be the caller's own. It must be safe to +// call from any thread, and must not call into the library. +typedef void (*WaifuModelManagerLogCallback)(void *user_data, const WaifuModelManagerLogEvent *event); + +typedef struct { + // 0..1 of the file, or -1 where the server did not say how long it is. + double fraction; + // Bytes of the file; total is 0 where it is not known. + uint64_t done; + uint64_t total; + // Which file of how many; parts is 0 until the first has named the rest. + uint32_t part; + uint32_t parts; + // The file, or NULL before the first. + const char *file; + // "fetching sdxl:base from Hugging Face -- ... (1 of 3)", for a status line. + const char *words; + // Since the job started. + double seconds; + // Which hub the files come from, once that is settled; NULL before. + const char *hub; +} WaifuModelManagerProgressEvent; + +typedef void (*WaifuModelManagerProgressCallback)(void *user_data, + const WaifuModelManagerProgressEvent *event); + +// How a job ended, as on_complete is told it. +typedef struct { + WaifuModelManagerStatusCode code; + // Why, when code is not WAIFU_MODELMANAGER_OK; NULL when it is. + const char *message; +} WaifuModelManagerStatus; + +// How a fetch ended. +typedef struct { + WaifuModelManagerStatus status; +} WaifuModelManagerCompletionEvent; + +typedef void (*WaifuModelManagerCompleteCallback)(void *user_data, + const WaifuModelManagerCompletionEvent *event); + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +// The ABI this library was built with: WAIFU_MODELMANAGER_ABI_VERSION. +uint32_t waifu_modelmanager_abi_version(void); + +// This thread's last message from this library; valid until this thread's next call into it. +const char *waifu_modelmanager_last_error(void); + +// Frees a string a call into this library returned. NULL is fine. +void waifu_modelmanager_string_free(char *s); + +// Every line this library would print goes to the callback instead, from whichever thread wrote +// it. NULL puts them back on the console. libwaifu's lines are waifu_set_log_callback's. +void waifu_modelmanager_set_log_callback(void *user_data, WaifuModelManagerLogCallback callback); + +// Lines below the level are not made at all. INFO unless set. +void waifu_modelmanager_set_log_level(WaifuModelManagerLogLevel level); + +// The published models, aliases only, as JSON. Free with waifu_modelmanager_string_free. NULL on +// failure. +char *waifu_modelmanager_catalog_json(void); + +// Where downloads go. Free with waifu_modelmanager_string_free. NULL on failure. +char *waifu_modelmanager_directory(void); + +// Where downloads go, who said so and what else could, as JSON: +// `{"directory", "said": "environment"|"config"|"default", "default", "config_file"}`, the last +// two null where they cannot be told. Free with waifu_modelmanager_string_free. NULL on failure. +char *waifu_modelmanager_directory_json(void); + +// The manifest of a published model, as a path for libwaifu's load and describe: the one in the +// download directory where it has been fetched, and otherwise one of the bundled manifests. Only +// a model whose catalogue entry says `cached` has its weights beside it. A manifest's own path is +// handed back as it is. Free with waifu_modelmanager_string_free. NULL, with +// waifu_modelmanager_last_error, where there is no manifest here: fetch the model, or every +// manifest, first. +char *waifu_modelmanager_manifest_path(const char *name); + +// What a screen should start a model's boxes at, as JSON: the kind of model's defaults with what +// its manifest's `suggested:` block says laid over them. `model` is the path of a manifest -- see +// waifu_modelmanager_manifest_path -- or "tones". A picture model's are `width`, `height`, `steps`, +// `guidance`, `takes_guidance`, `sizes`, `prompt` and `avoid`; a voice's are `speed` and +// `temperature`, and `conversion.steps` for one that also converts. None of it is needed to run a +// model, whose request says every number itself. Free with waifu_modelmanager_string_free. NULL, +// with waifu_modelmanager_last_error, for anything that is not a manifest on the disk. +char *waifu_modelmanager_defaults_json(const char *model); + +// Downloads go to `path` from now on, saved to config.toml; NULL or "" for the default. +WaifuModelManagerStatusCode waifu_modelmanager_set_directory(const char *path); + +// Reads the manifests of published models out of `directory` where the download directory has +// none: an app built with every model's manifest (waifu-tools' `fetch_manifests` fetches them) +// can describe a model before anything is downloaded: waifu_modelmanager_manifest_path hands +// back the bundled one where nothing has been fetched. Laid out as the download directory is, +// and only read. NULL or "" for none. Not saved. +WaifuModelManagerStatusCode waifu_modelmanager_set_bundled_manifests(const char *directory); + +// Deletes a downloaded package. +WaifuModelManagerStatusCode waifu_modelmanager_remove(const char *name); + +// Starts fetching a published model on a thread of its own and returns at once. Fetching one +// already here finishes at once. Where it went is waifu_modelmanager_manifest_path's to say. +// NULL, with waifu_modelmanager_last_error and no callback ever, where it cannot start. +WaifuModelFetch *waifu_modelmanager_fetch_async(const char *model, + void *user_data, + WaifuModelManagerProgressCallback on_progress, + WaifuModelManagerCompleteCallback on_complete); + +// Starts fetching the manifest of every published model that is not here yet -- a couple of +// kilobytes each, and none of their packages -- on a thread of its own, and returns at once. +// What a model suggests is in its manifest, so a screen that has called this can fill in a +// model's settings the moment it is chosen. Each manifest is a part of the whole in the progress; +// with every one already here it finishes at once. Cancelled and freed as a model's fetch is. +// NULL, with waifu_modelmanager_last_error and no callback ever, where it cannot start. +WaifuModelFetch *waifu_modelmanager_fetch_manifests_async(void *user_data, + WaifuModelManagerProgressCallback on_progress, + WaifuModelManagerCompleteCallback on_complete); + +void waifu_modelmanager_fetch_cancel(WaifuModelFetch *fetch); + +// After on_complete; cancels first if it has not had it, and waits for the thread to end. +void waifu_modelmanager_fetch_free(WaifuModelFetch *fetch); + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus + +#endif /* LIBWAIFU_WAIFU_MODELMANAGER_H */ diff --git a/waifu-modelmanager/src/c_api.rs b/waifu-modelmanager/src/c_api.rs new file mode 100644 index 0000000..5823b81 --- /dev/null +++ b/waifu-modelmanager/src/c_api.rs @@ -0,0 +1,417 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! The C API: the catalogue, fetching, where models are kept, and the log, as `extern "C"`. +//! Everything outside this crate goes through these. +//! +//! Nothing with any logic in it lives here. This converts -- C strings to `&str`, what comes back +//! to borrowed C views, panics to status codes -- and keeps the promises docs/ffi.md makes about +//! when a callback can arrive. + +#![allow(clippy::missing_safety_doc)] + +use std::ffi::{c_char, c_void}; +use std::panic::{self, AssertUnwindSafe}; +use std::ptr; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::thread::{self, JoinHandle}; +use std::time::Instant; + +mod support; + +pub use support::{ + WaifuModelManagerCompleteCallback, WaifuModelManagerCompletionEvent, + WaifuModelManagerLogCallback, WaifuModelManagerLogEvent, WaifuModelManagerLogLevel, + WaifuModelManagerProgressCallback, WaifuModelManagerProgressEvent, WaifuModelManagerStatus, + WaifuModelManagerStatusCode, +}; +use support::{ + c_string, caught, required, set_last_error, status, string_in, string_out, Gate, UserData, +}; + +use crate::fetch::{self, Failed, Fetch}; +use crate::{config, hub, log}; + +use WaifuModelManagerStatusCode::*; + +/// What a caller compiled against this header expects. Bumped when anything in it changes in a +/// way an older caller would misread. +pub const WAIFU_MODELMANAGER_ABI_VERSION: u32 = 1; + +/// The ABI this library was built with: WAIFU_MODELMANAGER_ABI_VERSION. +#[no_mangle] +pub extern "C" fn waifu_modelmanager_abi_version() -> u32 { + WAIFU_MODELMANAGER_ABI_VERSION +} + +/// This thread's last message from this library; valid until this thread's next call into it. +#[no_mangle] +pub extern "C" fn waifu_modelmanager_last_error() -> *const c_char { + support::last_error() +} + +/// Frees a string a call into this library returned. NULL is fine. +#[no_mangle] +pub unsafe extern "C" fn waifu_modelmanager_string_free(s: *mut c_char) { + support::free_string(s) +} + +// -- logging --------------------------------------------------------------------------------- + +/// Every line this library would print goes to the callback instead, from whichever thread wrote +/// it. NULL puts them back on the console. libwaifu's lines are waifu_set_log_callback's. +#[no_mangle] +pub extern "C" fn waifu_modelmanager_set_log_callback(user_data: *mut c_void, callback: WaifuModelManagerLogCallback) { + caught((), || log::set_sink(support::log_sink(user_data, callback))) +} + +/// Lines below the level are not made at all. INFO unless set. +#[no_mangle] +pub extern "C" fn waifu_modelmanager_set_log_level(level: WaifuModelManagerLogLevel) { + caught((), || log::set_level(support::log_level(level))) +} + +// -- the model manager ----------------------------------------------------------------------- + +/// The published models, aliases only, as JSON. Free with waifu_modelmanager_string_free. NULL on +/// failure. +#[no_mangle] +pub extern "C" fn waifu_modelmanager_catalog_json() -> *mut c_char { + caught(ptr::null_mut(), || string_out(&fetch::catalog_json().to_string())) +} + +/// Where downloads go. Free with waifu_modelmanager_string_free. NULL on failure. +#[no_mangle] +pub extern "C" fn waifu_modelmanager_directory() -> *mut c_char { + caught(ptr::null_mut(), || match hub::model_directory() { + Ok((directory, _)) => string_out(&directory.to_string_lossy()), + Err(error) => { + set_last_error(&error.to_string()); + ptr::null_mut() + } + }) +} + +/// Where downloads go, who said so and what else could, as JSON: +/// `{"directory", "said": "environment"|"config"|"default", "default", "config_file"}`, the last +/// two null where they cannot be told. Free with waifu_modelmanager_string_free. NULL on failure. +#[no_mangle] +pub extern "C" fn waifu_modelmanager_directory_json() -> *mut c_char { + caught(ptr::null_mut(), || match hub::model_directory() { + Ok((directory, said)) => string_out( + &serde_json::json!({ + "directory": directory.to_string_lossy(), + "said": match said { + hub::Said::Environment => "environment", + hub::Said::Config => "config", + hub::Said::Default => "default", + }, + "default": hub::default_directory().ok().map(|default| default.to_string_lossy().into_owned()), + "config_file": config::path().map(|file| file.to_string_lossy().into_owned()), + }) + .to_string(), + ), + Err(error) => { + set_last_error(&error.to_string()); + ptr::null_mut() + } + }) +} + +/// The manifest of a published model, as a path for libwaifu's load and describe: the one in the +/// download directory where it has been fetched, and otherwise one of the bundled manifests. Only +/// a model whose catalogue entry says `cached` has its weights beside it. A manifest's own path is +/// handed back as it is. Free with waifu_modelmanager_string_free. NULL, with +/// waifu_modelmanager_last_error, where there is no manifest here: fetch the model, or every +/// manifest, first. +#[no_mangle] +pub unsafe extern "C" fn waifu_modelmanager_manifest_path(name: *const c_char) -> *mut c_char { + caught(ptr::null_mut(), || { + let name = match required(name, "name") { + Ok(name) => name, + Err(error) => { + set_last_error(&error); + return ptr::null_mut(); + } + }; + if hub::full_name(&name).is_none() { + if std::path::Path::new(&name).is_file() { + return string_out(&name); + } + set_last_error(&format!("there is no model called \"{name}\": neither a published name nor a manifest that exists")); + return ptr::null_mut(); + } + match hub::local_manifest(&name) { + Some(path) => string_out(&path.to_string_lossy()), + None => { + set_last_error(&format!("the manifest of \"{name}\" is not here: fetch it first")); + ptr::null_mut() + } + } + }) +} + +/// What a screen should start a model's boxes at, as JSON: the kind of model's defaults with what +/// its manifest's `suggested:` block says laid over them. `model` is the path of a manifest -- see +/// waifu_modelmanager_manifest_path -- or "tones". A picture model's are `width`, `height`, `steps`, +/// `guidance`, `takes_guidance`, `sizes`, `prompt` and `avoid`; a voice's are `speed` and +/// `temperature`, and `conversion.steps` for one that also converts. None of it is needed to run a +/// model, whose request says every number itself. Free with waifu_modelmanager_string_free. NULL, +/// with waifu_modelmanager_last_error, for anything that is not a manifest on the disk. +#[no_mangle] +pub unsafe extern "C" fn waifu_modelmanager_defaults_json(model: *const c_char) -> *mut c_char { + caught(ptr::null_mut(), || { + let model = match required(model, "model") { + Ok(model) => model, + Err(error) => { + set_last_error(&error); + return ptr::null_mut(); + } + }; + match crate::defaults::defaults_json(&model) { + Some(defaults) => string_out(&defaults.to_string()), + None => { + set_last_error(&format!("there is no manifest at \"{model}\"")); + ptr::null_mut() + } + } + }) +} + +/// Downloads go to `path` from now on, saved to config.toml; NULL or "" for the default. +#[no_mangle] +pub unsafe extern "C" fn waifu_modelmanager_set_directory(path: *const c_char) -> WaifuModelManagerStatusCode { + caught(WAIFU_MODELMANAGER_ERR_INTERNAL, || { + let path = match string_in(path, "path") { + Ok(path) => path.filter(|path| !path.is_empty()), + Err(error) => { + set_last_error(&error); + return WAIFU_MODELMANAGER_ERR_INVALID_ARGUMENT; + } + }; + match config::set_model_dir(path.as_deref().map(std::path::Path::new)) { + Ok(_) => WAIFU_MODELMANAGER_OK, + Err(error) => { + set_last_error(&error.to_string()); + WAIFU_MODELMANAGER_ERR_INTERNAL + } + } + }) +} + +/// Reads the manifests of published models out of `directory` where the download directory has +/// none: an app built with every model's manifest (waifu-tools' `fetch_manifests` fetches them) +/// can describe a model before anything is downloaded: waifu_modelmanager_manifest_path hands +/// back the bundled one where nothing has been fetched. Laid out as the download directory is, +/// and only read. NULL or "" for none. Not saved. +#[no_mangle] +pub unsafe extern "C" fn waifu_modelmanager_set_bundled_manifests(directory: *const c_char) -> WaifuModelManagerStatusCode { + caught(WAIFU_MODELMANAGER_ERR_INTERNAL, || { + let directory = match string_in(directory, "directory") { + Ok(directory) => directory.filter(|directory| !directory.is_empty()), + Err(error) => { + set_last_error(&error); + return WAIFU_MODELMANAGER_ERR_INVALID_ARGUMENT; + } + }; + hub::set_bundled_manifests(directory.map(std::path::PathBuf::from)); + WAIFU_MODELMANAGER_OK + }) +} + +/// Deletes a downloaded package. +#[no_mangle] +pub unsafe extern "C" fn waifu_modelmanager_remove(name: *const c_char) -> WaifuModelManagerStatusCode { + caught(WAIFU_MODELMANAGER_ERR_INTERNAL, || { + let name = match required(name, "name") { + Ok(name) => name, + Err(error) => { + set_last_error(&error); + return WAIFU_MODELMANAGER_ERR_INVALID_ARGUMENT; + } + }; + if hub::full_name(&name).is_none() { + set_last_error(&format!("there is no published model called \"{name}\"")); + return WAIFU_MODELMANAGER_ERR_UNKNOWN_MODEL; + } + match hub::remove(&name) { + Ok(()) => WAIFU_MODELMANAGER_OK, + Err(error) => { + set_last_error(&error.to_string()); + WAIFU_MODELMANAGER_ERR_INTERNAL + } + } + }) +} + +/// A download running on a thread of its own. +pub struct WaifuModelFetch { + stop: Arc, + thread: Option>, +} + +/// Starts fetching a published model on a thread of its own and returns at once. Fetching one +/// already here finishes at once. Where it went is waifu_modelmanager_manifest_path's to say. +/// NULL, with waifu_modelmanager_last_error and no callback ever, where it cannot start. +#[no_mangle] +pub unsafe extern "C" fn waifu_modelmanager_fetch_async( + model: *const c_char, + user_data: *mut c_void, + on_progress: WaifuModelManagerProgressCallback, + on_complete: WaifuModelManagerCompleteCallback, +) -> *mut WaifuModelFetch { + caught(ptr::null_mut(), || { + let model = match required(model, "model") { + Ok(model) => model, + Err(error) => { + set_last_error(&error); + return ptr::null_mut(); + } + }; + spawn_fetch(user_data, on_progress, on_complete, move |on_progress, stop| { + fetch::fetch(&model, on_progress, stop).map(|_| ()) + }) + }) +} + +/// Starts fetching the manifest of every published model that is not here yet -- a couple of +/// kilobytes each, and none of their packages -- on a thread of its own, and returns at once. +/// What a model suggests is in its manifest, so a screen that has called this can fill in a +/// model's settings the moment it is chosen. Each manifest is a part of the whole in the progress; +/// with every one already here it finishes at once. Cancelled and freed as a model's fetch is. +/// NULL, with waifu_modelmanager_last_error and no callback ever, where it cannot start. +#[no_mangle] +pub unsafe extern "C" fn waifu_modelmanager_fetch_manifests_async( + user_data: *mut c_void, + on_progress: WaifuModelManagerProgressCallback, + on_complete: WaifuModelManagerCompleteCallback, +) -> *mut WaifuModelFetch { + caught(ptr::null_mut(), || { + spawn_fetch(user_data, on_progress, on_complete, fetch::fetch_manifests) + }) +} + +/// Runs `work` on a thread of its own, with the callbacks it reports to and a stop it can be +/// cancelled by, and hands back the handle that stop is reached through. +fn spawn_fetch( + user_data: *mut c_void, + on_progress: WaifuModelManagerProgressCallback, + on_complete: WaifuModelManagerCompleteCallback, + work: impl FnOnce(&mut dyn FnMut(&Fetch), &dyn Fn() -> bool) -> Result<(), Failed> + Send + 'static, +) -> *mut WaifuModelFetch { + let gate = Gate::default(); + let user_data = UserData(user_data); + let stop = Arc::new(AtomicBool::new(false)); + + let thread = thread::Builder::new().name("waifu-fetch".to_string()).spawn({ + let stop = Arc::clone(&stop); + let gate = gate.clone(); + move || { + let started = Instant::now(); + let mut tell = |fetch: &Fetch| { + if on_progress.is_none() { + return; + } + gate.wait(); + let words = fetch.words(); + support::tell_progress( + on_progress, + user_data, + &support::Progress { + fraction: fetch.fraction(), + done: fetch.done, + total: fetch.total.unwrap_or(0), + part: fetch.part as u32, + parts: fetch.parts as u32, + file: Some(fetch.file.as_str()).filter(|file| !file.is_empty()), + words: &words, + seconds: started.elapsed().as_secs_f64(), + hub: fetch.hub, + }, + ); + }; + let ended = panic::catch_unwind(AssertUnwindSafe(|| { + work(&mut tell, &|| stop.load(Ordering::Relaxed)) + })) + .unwrap_or_else(|_| Err(Failed::Fetch("a bug in the library: the fetch panicked".to_string()))); + + if on_complete.is_none() { + return; + } + gate.wait(); + let (code, message) = match &ended { + Ok(()) => (WAIFU_MODELMANAGER_OK, None), + Err(Failed::Stopped) => (WAIFU_MODELMANAGER_CANCELLED, None), + Err(Failed::UnknownModel(message)) => (WAIFU_MODELMANAGER_ERR_UNKNOWN_MODEL, Some(message.as_str())), + Err(Failed::Fetch(message)) => (WAIFU_MODELMANAGER_ERR_FETCH, Some(message.as_str())), + }; + tell_completion(on_complete, user_data, code, message); + } + }); + match thread { + Ok(thread) => { + let fetch = Box::into_raw(Box::new(WaifuModelFetch { + stop, + thread: Some(thread), + })); + gate.open(); + fetch + } + Err(error) => { + set_last_error(&format!("could not start the download's thread: {error}")); + ptr::null_mut() + } + } +} + +/// Calls `callback`, if there is one, with how a job that makes nothing ended. +pub fn tell_completion(callback: WaifuModelManagerCompleteCallback, user_data: UserData, code: WaifuModelManagerStatusCode, message: Option<&str>) { + let Some(callback) = callback else { + return; + }; + let message = message.map(c_string); + let event = WaifuModelManagerCompletionEvent { + status: status(code, &message), + }; + unsafe { callback(user_data.0, &event) }; +} + +#[no_mangle] +pub unsafe extern "C" fn waifu_modelmanager_fetch_cancel(fetch: *mut WaifuModelFetch) { + if let Some(fetch) = fetch.as_ref() { + fetch.stop.store(true, Ordering::Relaxed); + } +} + +/// After on_complete; cancels first if it has not had it, and waits for the thread to end. +#[no_mangle] +pub unsafe extern "C" fn waifu_modelmanager_fetch_free(fetch: *mut WaifuModelFetch) { + if fetch.is_null() { + return; + } + let mut fetch = Box::from_raw(fetch); + fetch.stop.store(true, Ordering::Relaxed); + if let Some(thread) = fetch.thread.take() { + let _ = thread.join(); + } +} + +#[cfg(test)] +mod tests; diff --git a/waifu-modelmanager/src/c_api/support.rs b/waifu-modelmanager/src/c_api/support.rs new file mode 100644 index 0000000..5234380 --- /dev/null +++ b/waifu-modelmanager/src/c_api/support.rs @@ -0,0 +1,289 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! What the C API is made of besides its functions: the types every part of it takes and hands +//! back -- status, the log, progress and completion -- and the helpers that convert to and from +//! them. Written into waifu_modelmanager.h with the functions, under this library's own names: +//! libwaifu has types of its own, and the two headers share nothing. + +#![allow(non_camel_case_types)] +#![allow(clippy::missing_safety_doc)] + +use std::cell::RefCell; +use std::ffi::{c_char, c_void, CStr, CString}; +use std::panic::{self, AssertUnwindSafe}; +use std::ptr; +use std::sync::{Arc, Condvar, Mutex}; + +use crate::log::{Level, Sink}; + +// -- status ---------------------------------------------------------------------------------- + +/// What a call or a job came to. +#[repr(C)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum WaifuModelManagerStatusCode { + WAIFU_MODELMANAGER_OK = 0, + WAIFU_MODELMANAGER_CANCELLED = 1, + /// A NULL, or a string that is not UTF-8. + WAIFU_MODELMANAGER_ERR_INVALID_ARGUMENT = 2, + /// No published model by that name, or no manifest at that path. + WAIFU_MODELMANAGER_ERR_UNKNOWN_MODEL = 3, + /// The download failed. + WAIFU_MODELMANAGER_ERR_FETCH = 4, + /// A panic in the library, or a file that could not be read or written. + WAIFU_MODELMANAGER_ERR_INTERNAL = 5, +} + +/// How a job ended, as on_complete is told it. +#[repr(C)] +pub struct WaifuModelManagerStatus { + pub code: WaifuModelManagerStatusCode, + /// Why, when code is not WAIFU_MODELMANAGER_OK; NULL when it is. + pub message: *const c_char, +} + +thread_local! { + static LAST_ERROR: RefCell = RefCell::new(CString::default()); +} + +/// Says why this thread's last call failed, for its `_last_error` to hand back. +pub fn set_last_error(message: &str) { + let message = c_string(message); + LAST_ERROR.with(|last| *last.borrow_mut() = message); +} + +/// This thread's last message; valid until this thread's next call. +pub fn last_error() -> *const c_char { + LAST_ERROR.with(|last| last.borrow().as_ptr()) +} + +/// Frees a string [`string_out`] made. NULL is fine. +pub unsafe fn free_string(s: *mut c_char) { + if !s.is_null() { + drop(CString::from_raw(s)); + } +} + +/// Runs a synchronous entry point, turning a panic into its message in the last error and +/// `otherwise` as what the call returns. +pub fn caught(otherwise: T, body: impl FnOnce() -> T) -> T { + match panic::catch_unwind(AssertUnwindSafe(body)) { + Ok(value) => value, + Err(cause) => { + let said = cause + .downcast_ref::<&str>() + .map(|said| said.to_string()) + .or_else(|| cause.downcast_ref::().cloned()) + .unwrap_or_else(|| "it panicked".to_string()); + set_last_error(&format!("a bug in the library: {said}")); + otherwise + } + } +} + +/// A string handed in, or the reason it cannot be read. NULL is the caller's to decide. +pub unsafe fn string_in(pointer: *const c_char, what: &str) -> Result, String> { + if pointer.is_null() { + return Ok(None); + } + CStr::from_ptr(pointer) + .to_str() + .map(|text| Some(text.to_string())) + .map_err(|_| format!("{what} is not UTF-8")) +} + +/// A string handed in that may not be NULL. +pub unsafe fn required(pointer: *const c_char, what: &str) -> Result { + string_in(pointer, what)?.ok_or_else(|| format!("{what} is NULL")) +} + +/// A string for the caller to free with the library's `_string_free`. +pub fn string_out(text: &str) -> *mut c_char { + c_string(text).into_raw() +} + +/// `text` as C can read it: a NUL inside it would end it early, so it is a space instead. +pub fn c_string(text: &str) -> CString { + CString::new(text.replace('\0', " ")).unwrap_or_default() +} + +// -- logging --------------------------------------------------------------------------------- + +#[repr(C)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum WaifuModelManagerLogLevel { + WAIFU_MODELMANAGER_LOG_DEBUG = 0, + WAIFU_MODELMANAGER_LOG_INFO = 1, + WAIFU_MODELMANAGER_LOG_WARNING = 2, + WAIFU_MODELMANAGER_LOG_ERROR = 3, + /// The last line: the process aborts when the callback returns. + WAIFU_MODELMANAGER_LOG_FATAL = 4, +} + +#[repr(C)] +pub struct WaifuModelManagerLogEvent { + pub level: WaifuModelManagerLogLevel, + /// Where it was written: a module's name. + pub source: *const c_char, + /// The line alone, with no level or time in front of it. + pub message: *const c_char, +} + +/// Called on whichever thread wrote the line, which can be the caller's own. It must be safe to +/// call from any thread, and must not call into the library. +pub type WaifuModelManagerLogCallback = Option; + +/// A caller's pointer, handed back as it was given, on whichever thread a callback runs. +#[derive(Clone, Copy)] +pub struct UserData(pub *mut c_void); +unsafe impl Send for UserData {} +unsafe impl Sync for UserData {} + +/// The sink that calls `callback`, or None -- the console -- where there is no callback. +pub fn log_sink(user_data: *mut c_void, callback: WaifuModelManagerLogCallback) -> Option { + let callback = callback?; + let user_data = UserData(user_data); + Some(Box::new(move |level, source, message| { + let source = c_string(source); + let message = c_string(message); + let event = WaifuModelManagerLogEvent { + level: match level { + Level::Debug => WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_DEBUG, + Level::Info => WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_INFO, + Level::Warning => WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_WARNING, + Level::Error => WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_ERROR, + Level::Fatal => WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_FATAL, + }, + source: source.as_ptr(), + message: message.as_ptr(), + }; + let user_data = user_data; + unsafe { callback(user_data.0, &event) }; + })) +} + +pub fn log_level(level: WaifuModelManagerLogLevel) -> Level { + match level { + WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_DEBUG => Level::Debug, + WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_INFO => Level::Info, + WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_WARNING => Level::Warning, + WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_ERROR => Level::Error, + WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_FATAL => Level::Fatal, + } +} + +// -- progress -------------------------------------------------------------------------------- + +#[repr(C)] +pub struct WaifuModelManagerProgressEvent { + /// 0..1 of the file, or -1 where the server did not say how long it is. + pub fraction: f64, + /// Bytes of the file; total is 0 where it is not known. + pub done: u64, + pub total: u64, + /// Which file of how many; parts is 0 until the first has named the rest. + pub part: u32, + pub parts: u32, + /// The file, or NULL before the first. + pub file: *const c_char, + /// "fetching sdxl:base from Hugging Face -- ... (1 of 3)", for a status line. + pub words: *const c_char, + /// Since the job started. + pub seconds: f64, + /// Which hub the files come from, once that is settled; NULL before. + pub hub: *const c_char, +} + +pub type WaifuModelManagerProgressCallback = Option; + +/// How a fetch ended. +#[repr(C)] +pub struct WaifuModelManagerCompletionEvent { + pub status: WaifuModelManagerStatus, +} + +pub type WaifuModelManagerCompleteCallback = Option; + +/// A progress report, before it is made into C. +pub struct Progress<'a> { + pub fraction: Option, + pub done: u64, + pub total: u64, + pub part: u32, + pub parts: u32, + pub file: Option<&'a str>, + pub words: &'a str, + pub seconds: f64, + pub hub: Option<&'a str>, +} + +/// Calls `callback`, if there is one, with `progress`. +pub fn tell_progress(callback: WaifuModelManagerProgressCallback, user_data: UserData, progress: &Progress) { + let Some(callback) = callback else { + return; + }; + let file = progress.file.map(c_string); + let words = c_string(progress.words); + let hub = progress.hub.map(c_string); + let event = WaifuModelManagerProgressEvent { + fraction: progress.fraction.unwrap_or(-1.0), + done: progress.done, + total: progress.total, + part: progress.part, + parts: progress.parts, + file: file.as_ref().map_or(ptr::null(), |file| file.as_ptr()), + words: words.as_ptr(), + seconds: progress.seconds, + hub: hub.as_ref().map_or(ptr::null(), |hub| hub.as_ptr()), + }; + unsafe { callback(user_data.0, &event) }; +} + +/// A status whose message lives as long as `message` does. +pub fn status(code: WaifuModelManagerStatusCode, message: &Option) -> WaifuModelManagerStatus { + WaifuModelManagerStatus { + code, + message: message.as_ref().map_or(ptr::null(), |message| message.as_ptr()), + } +} + +// -- keeping the promises about when a callback can arrive --------------------------------------- + +/// Shut until the `_async` call that made it has returned: a callback that arrives before then +/// waits. A job's thread can finish before the call that started it is done handing back its +/// handle, and a caller whose continuation is set up from that handle would miss the end. +#[derive(Clone, Default)] +pub struct Gate(Arc<(Mutex, Condvar)>); + +impl Gate { + pub fn open(&self) { + let (open, opened) = &*self.0; + *open.lock().unwrap_or_else(|held| held.into_inner()) = true; + opened.notify_all(); + } + + pub fn wait(&self) { + let (open, opened) = &*self.0; + let mut is_open = open.lock().unwrap_or_else(|held| held.into_inner()); + while !*is_open { + is_open = opened.wait(is_open).unwrap_or_else(|held| held.into_inner()); + } + } +} diff --git a/waifu-modelmanager/src/c_api/tests.rs b/waifu-modelmanager/src/c_api/tests.rs new file mode 100644 index 0000000..38e168b --- /dev/null +++ b/waifu-modelmanager/src/c_api/tests.rs @@ -0,0 +1,187 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! The C API, called the way C calls it: through the `extern "C"` functions, with callbacks that +//! are `extern "C"` functions and a `user_data` pointer -- here, to a channel the test waits on. +//! Nothing here touches the network. + +use super::*; + +use std::ffi::CStr; +use std::sync::mpsc::{self, Receiver, Sender}; +use std::sync::Mutex; +use std::time::Duration; + + +const WAIT: Duration = Duration::from_secs(20); + +/// What a test hears of one fetch, sent from the callbacks. +#[derive(Debug, PartialEq)] +enum Heard { + Progress(String), + Complete(WaifuModelManagerStatusCode, Option), +} + +unsafe fn text(pointer: *const c_char) -> Option { + (!pointer.is_null()).then(|| CStr::from_ptr(pointer).to_string_lossy().into_owned()) +} + +unsafe fn sender<'a>(user_data: *mut c_void) -> &'a Sender { + &*(user_data as *const Sender) +} + +unsafe extern "C" fn on_progress(user_data: *mut c_void, event: *const WaifuModelManagerProgressEvent) { + let event = &*event; + let _ = sender(user_data).send(Heard::Progress(text(event.words).unwrap_or_default())); +} + +unsafe extern "C" fn on_complete(user_data: *mut c_void, event: *const WaifuModelManagerCompletionEvent) { + let status = &(*event).status; + let _ = sender(user_data).send(Heard::Complete(status.code, text(status.message))); +} + +fn last_error() -> String { + unsafe { text(waifu_modelmanager_last_error()) }.unwrap_or_default() +} + +#[test] +fn the_abi_is_the_header_s() { + assert_eq!(waifu_modelmanager_abi_version(), WAIFU_MODELMANAGER_ABI_VERSION); +} + +#[test] +fn what_is_published_is_json_and_freed_by_the_caller() { + unsafe { + let catalog = waifu_modelmanager_catalog_json(); + assert!(!catalog.is_null()); + assert!(text(catalog).unwrap().contains("\"sdxl:base\"")); + waifu_modelmanager_string_free(catalog); + } +} + +#[test] +fn where_models_go_is_json_and_says_who_said_so() { + unsafe { + let said = waifu_modelmanager_directory_json(); + assert!(!said.is_null(), "{}", last_error()); + let text = text(said).unwrap(); + waifu_modelmanager_string_free(said); + for key in ["\"directory\"", "\"said\"", "\"default\"", "\"config_file\""] { + assert!(text.contains(key), "{text}"); + } + } +} + +#[test] +fn a_manifest_path_is_a_path_or_none() { + unsafe { + // A path that is a manifest is handed back as it is. + let file = std::env::temp_dir().join(format!("waifu-c-api-{}.yaml", std::process::id())); + std::fs::write(&file, "weights:\n - a.safetensors\n").unwrap(); + let asked = std::ffi::CString::new(file.to_str().unwrap()).unwrap(); + let path = waifu_modelmanager_manifest_path(asked.as_ptr()); + assert_eq!(text(path).as_deref(), file.to_str()); + waifu_modelmanager_string_free(path); + std::fs::remove_file(&file).unwrap(); + + assert!(waifu_modelmanager_manifest_path(c"no-such-model:v9".as_ptr()).is_null()); + assert!(last_error().contains("no-such-model:v9"), "{}", last_error()); + assert!(waifu_modelmanager_manifest_path(ptr::null()).is_null()); + } +} + +#[test] +fn bundled_manifests_are_a_directory_or_none() { + unsafe { + // Where they are is the caller's; what is read from them is hub's to test. + let directory = std::env::temp_dir(); + let directory = std::ffi::CString::new(directory.to_str().expect("a UTF-8 path")).unwrap(); + assert_eq!(waifu_modelmanager_set_bundled_manifests(directory.as_ptr()), WAIFU_MODELMANAGER_OK); + assert_eq!(waifu_modelmanager_set_bundled_manifests(c"".as_ptr()), WAIFU_MODELMANAGER_OK); + assert_eq!(waifu_modelmanager_set_bundled_manifests(ptr::null()), WAIFU_MODELMANAGER_OK); + } +} + +#[test] +fn removing_a_model_nobody_published_is_refused_by_name() { + unsafe { + assert_eq!(waifu_modelmanager_remove(c"no-such-model:v9".as_ptr()), WAIFU_MODELMANAGER_ERR_UNKNOWN_MODEL); + assert!(last_error().contains("no-such-model:v9"), "{}", last_error()); + assert_eq!(waifu_modelmanager_remove(ptr::null()), WAIFU_MODELMANAGER_ERR_INVALID_ARGUMENT); + } +} + +#[test] +fn a_fetch_refused_at_the_door_is_null_and_calls_nothing() { + unsafe { + let (send, heard) = mpsc::channel(); + let user_data = &send as *const Sender as *mut c_void; + let fetch = waifu_modelmanager_fetch_async(ptr::null(), user_data, Some(on_progress), Some(on_complete)); + assert!(fetch.is_null()); + assert!(last_error().contains("model"), "{}", last_error()); + assert!(heard.recv_timeout(Duration::from_millis(100)).is_err()); + } +} + +#[test] +fn a_model_nobody_published_is_unknown_at_the_end() { + unsafe { + let (send, heard): (Sender, Receiver) = mpsc::channel(); + let user_data = &send as *const Sender as *mut c_void; + let fetch = waifu_modelmanager_fetch_async( + c"no-such-model:v9".as_ptr(), + user_data, + Some(on_progress), + Some(on_complete), + ); + assert!(!fetch.is_null(), "{}", last_error()); + match heard.recv_timeout(WAIT).expect("an end") { + Heard::Complete(WAIFU_MODELMANAGER_ERR_UNKNOWN_MODEL, Some(message)) => { + assert!(message.contains("no-such-model:v9"), "{message}") + } + other => panic!("{other:?}"), + } + waifu_modelmanager_fetch_free(fetch); + } +} + +unsafe extern "C" fn on_log(user_data: *mut c_void, event: *const WaifuModelManagerLogEvent) { + let lines = &*(user_data as *const Mutex>); + let event = &*event; + lines + .lock() + .unwrap() + .push((event.level, text(event.message).unwrap_or_default())); +} + +#[test] +fn a_log_callback_hears_the_library_s_lines() { + let lines: Box>> = Box::default(); + waifu_modelmanager_set_log_callback(&*lines as *const _ as *mut c_void, Some(on_log)); + log::write(log::Level::Warning, "hub", "a line for the app"); + waifu_modelmanager_set_log_callback(ptr::null_mut(), None); + + let lines = lines.lock().unwrap(); + assert!( + lines + .iter() + .any(|(level, line)| *level == WaifuModelManagerLogLevel::WAIFU_MODELMANAGER_LOG_WARNING && line == "a line for the app"), + "{lines:?}" + ); +} diff --git a/waifu/src/config.rs b/waifu-modelmanager/src/config.rs similarity index 100% rename from waifu/src/config.rs rename to waifu-modelmanager/src/config.rs diff --git a/waifu-modelmanager/src/defaults.rs b/waifu-modelmanager/src/defaults.rs new file mode 100644 index 0000000..435e70a --- /dev/null +++ b/waifu-modelmanager/src/defaults.rs @@ -0,0 +1,252 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! What a screen should start a model's boxes at: the kind of model's own defaults, with what its +//! manifest's `suggested:` block says laid over them. +//! +//! None of it is needed to run a model -- a request says every number itself -- so it is the +//! model manager's to say, not the engine's. What the engine says of a model is what a request +//! may hold: `waifu_describe_json`. + +use std::collections::BTreeMap; +use std::path::Path; + +use serde_json::{json, Value}; + +use crate::suggested::Suggestions; +use crate::yaml::{self, Node}; + +/// The stand-in voice, which is libwaifu's own and has no manifest. +pub const TONES: &str = "tones"; + +/// The sizes offered for a model that does not name its own, which are the ones SDXL was trained +/// at. Every one of them is a multiple of the 32 pixels the U-Net's own halvings need. +pub const SIZES: &[(i32, i32)] = &[ + (512, 512), + (640, 640), + (768, 768), + (896, 896), + (1024, 1024), + (832, 1216), + (1216, 832), + (768, 1344), + (1344, 768), +]; + +/// How many sizes a model has to name before its list is used instead of the one above. +/// +/// One size is not a choice and two is barely one. A model that names fewer has not replaced the +/// list, it has emptied it, and the built-in one is the better thing to offer. +pub const ENOUGH_SIZES: usize = 3; + +/// The numbers a picture model's boxes start at. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Picture { + pub width: i32, + pub height: i32, + pub steps: i32, + pub guidance: f32, +} + +/// What a kind of model starts at before its manifest says anything, and whether guidance is a +/// thing to offer it at all -- false for Krea 2, whose only release here is distilled and has no +/// second pass. Anything that is not one of the other kinds is SDXL, which is what the engine +/// builds it as too. +fn kind_of_picture(kind: &str) -> (Picture, bool) { + let picture = |steps, guidance| Picture { + width: 1024, + height: 1024, + steps, + guidance, + }; + match kind { + "anima" => (picture(8, 1.0), true), + "krea2" => (picture(8, 1.0), false), + // Sampled without guidance by default, and able to take it: the reference's + // `true_cfg_scale` is a second pass on a negative prompt. + "qwen_image" => (picture(40, 1.0), true), + _ => (picture(30, 5.0), true), + } +} + +/// Whether to offer guidance and a negative prompt, for a package suggesting `suggested`, of a kind +/// that offers them where `kind_does`. +/// +/// What the package says, where it says it. Where it does not, a card that asks for a guidance of +/// one has said it all the same: at one there is no unconditional pass, so a dial above it is a +/// setting the card advises against and a negative prompt is text nothing reads. Anima Turbo is +/// the one published like that -- "at CFG 1 and 8-12 steps", and no `takes_guidance` key. Silence +/// with any other number leaves the kind's answer alone. +fn guided_by(suggested: &Suggestions, kind_does: bool) -> bool { + suggested + .takes_guidance + .unwrap_or(kind_does && suggested.guidance != Some(1.0)) +} + +/// A float as a box can show it: four places, which is finer than any setting steps. +fn showable(value: f32) -> f64 { + (f64::from(value) * 10_000.0).round() / 10_000.0 +} + +/// The top level of the manifest at `path`, where it is one a model could be built from: a +/// `weights` list with something in it and a `config` block. None for anything else, which is +/// described as what a model with nothing to say is. +fn manifest(path: &Path) -> Option> { + let text = std::fs::read_to_string(path).ok()?; + let top = yaml::parse(&text).ok()?.into_map("the manifest").ok()?; + let weights = top.get("weights")?.as_seq()?; + if weights.is_empty() || !matches!(top.get("config"), Some(Node::Map(_))) { + return None; + } + Some(top) +} + +/// `config: model: type:`, the kind of model a manifest is. +fn kind(top: &BTreeMap) -> Option { + let Some(Node::Map(config)) = top.get("config") else { + return None; + }; + let Some(Node::Map(model)) = config.get("model") else { + return None; + }; + model.get("type")?.as_str().map(str::to_string) +} + +/// The `suggested:` block, or nothing suggested. +fn suggestions(top: &BTreeMap) -> Suggestions { + match top.get("suggested") { + Some(Node::Map(fields)) => Suggestions::from_fields(fields), + _ => Suggestions::default(), + } +} + +/// What the boxes start at for the model `asked` names -- the path of a manifest, or "tones" -- +/// as `waifu_modelmanager_defaults_json` hands it out: +/// +/// - a picture model: `width`, `height`, `steps`, `guidance`, `takes_guidance`, `sizes`, `prompt` +/// and `avoid`; +/// - a voice: `speed` and `temperature`, and a `conversion` holding `steps` for one that also +/// converts. +/// +/// None for anything that is not a manifest on the disk. +pub fn defaults_json(asked: &str) -> Option { + if asked == TONES { + return Some(json!({ "speed": 1.0, "temperature": 0.8 })); + } + if !Path::new(asked).is_file() { + return None; + } + + let top = manifest(Path::new(asked)); + let kind = top.as_ref().and_then(kind).unwrap_or_default(); + match kind.as_str() { + "indextts2" => return Some(json!({ "speed": 1.0, "temperature": 0.8 })), + "cosyvoice3" => { + return Some(json!({ + "speed": 1.0, + "temperature": 1.0, + "conversion": { "steps": 10 }, + })) + } + _ => {} + } + + let suggested = top.as_ref().map(suggestions).unwrap_or_default(); + let (picture, kind_does) = kind_of_picture(&kind); + let picture = suggested.over(picture); + let sizes = match suggested.sizes.len() >= ENOUGH_SIZES { + true => suggested.sizes.clone(), + false => SIZES.to_vec(), + }; + Some(json!({ + "width": picture.width, + "height": picture.height, + "steps": picture.steps, + "guidance": showable(picture.guidance), + "takes_guidance": guided_by(&suggested, kind_does), + "sizes": sizes.iter().map(|(width, height)| json!([width, height])).collect::>(), + "prompt": suggested.prompt, + "avoid": suggested.avoid, + })) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_card_that_asks_for_a_guidance_of_one_takes_none() { + let turbo = Suggestions { + guidance: Some(1.0), + ..Suggestions::default() + }; + assert!(!guided_by(&turbo, true), "Anima Turbo's card says CFG 1 and nothing else"); + + // What the package says outranks what it suggests: an undistilled model that likes one. + let says_so = Suggestions { + guidance: Some(1.0), + takes_guidance: Some(true), + ..Suggestions::default() + }; + assert!(guided_by(&says_so, true)); + + // Silence with another number, or with none, is the kind's answer. + let anima = Suggestions { + guidance: Some(4.0), + ..Suggestions::default() + }; + assert!(guided_by(&anima, true)); + assert!(guided_by(&Suggestions::default(), true)); + assert!(!guided_by(&Suggestions::default(), false)); + } + + /// A manifest of `kind` with `suggested` under it, in a file of its own. + fn a_manifest(stem: &str, kind: &str, suggested: &str) -> String { + let path = std::env::temp_dir().join(format!("waifu-defaults-{}-{stem}.yaml", std::process::id())); + let text = format!("weights:\n - w.safetensors\nconfig:\n model:\n type: {kind}\n{suggested}"); + std::fs::write(&path, text).unwrap(); + path.to_str().unwrap().to_string() + } + + #[test] + fn a_turbo_release_starts_at_its_own_numbers_and_offers_no_guidance() { + let turbo = defaults_json(&a_manifest("turbo", "anima", "suggested:\n steps: 8\n guidance: 1\n")).unwrap(); + assert_eq!(turbo["steps"], 8); + assert_eq!(turbo["guidance"], 1.0); + assert_eq!(turbo["takes_guidance"], false); + + let sdxl = defaults_json(&a_manifest("sdxl", "sdxl", "")).unwrap(); + assert_eq!(sdxl["steps"], 30); + assert_eq!(sdxl["guidance"], 5.0); + assert_eq!(sdxl["takes_guidance"], true); + assert_eq!(sdxl["sizes"].as_array().unwrap().len(), SIZES.len()); + + let krea2 = defaults_json(&a_manifest("krea2", "krea2", "")).unwrap(); + assert_eq!(krea2["takes_guidance"], false); + } + + #[test] + fn a_voice_starts_at_its_kind_s_numbers() { + assert_eq!(defaults_json(TONES).unwrap()["temperature"], 0.8); + let cosy = defaults_json(&a_manifest("cosy", "cosyvoice3", "")).unwrap(); + assert_eq!(cosy["temperature"], 1.0); + assert_eq!(cosy["conversion"]["steps"], 10); + assert!(defaults_json("/no/such/model.yaml").is_none()); + } +} diff --git a/waifu-modelmanager/src/fetch.rs b/waifu-modelmanager/src/fetch.rs new file mode 100644 index 0000000..36cebe0 --- /dev/null +++ b/waifu-modelmanager/src/fetch.rs @@ -0,0 +1,256 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! Fetching a model, or every model's manifest, as one job that reports how far along it is: what +//! the C API runs on a thread of its own, and what the engine's load starts with. + +use std::path::{Path, PathBuf}; + +use crate::hub; + +/// Which package of a model is being fetched, and how far into it. +#[derive(Clone, Debug, PartialEq)] +pub struct Fetch { + pub model: String, + /// Which hub it is coming from, once the fetch has settled that. Unknown for the moment + /// before: it is decided by a probe, and a screen that guessed would be naming a hub while + /// the probe is still out. + pub hub: Option<&'static str>, + pub file: String, + pub done: u64, + pub total: Option, + pub part: usize, + /// How many packages there are in all, or zero while that is still unknown. + pub parts: usize, +} + +impl Fetch { + /// A fetch of `model` that has not heard anything yet. + pub fn of(model: impl Into) -> Fetch { + Fetch { + model: model.into(), + hub: None, + file: String::new(), + done: 0, + total: None, + part: 0, + parts: 0, + } + } + + /// Takes in what the hub said. + pub fn heard(&mut self, progress: hub::Progress) { + match progress { + hub::Progress::From { hub } => self.hub = Some(hub), + hub::Progress::Fetching { + file, + done, + total, + part, + parts, + } => { + self.file = file.to_string(); + (self.done, self.total, self.part, self.parts) = (done, total, part, parts); + } + hub::Progress::Fetched { + file, + bytes, + part, + parts, + } => { + self.file = file.to_string(); + (self.done, self.total, self.part, self.parts) = (bytes, Some(bytes), part, parts); + } + } + } + + /// How far into the file, between nothing and all of it, or `None` where the server did not + /// say how long it is -- which is also what the moment before the first byte looks like. + pub fn fraction(&self) -> Option { + self.total + .filter(|total| *total > 0) + .map(|total| self.done as f64 / total as f64) + } + + /// What it is busy with, in the words a bar carries. They say which model, not only which + /// file -- the file names in a package are not something anybody recognises -- and which hub, + /// once that is settled. + pub fn words(&self) -> String { + if self.file.is_empty() { + return format!("fetching {}", self.model); + } + let from = self.hub.map(|hub| format!(" from {hub}")).unwrap_or_default(); + match self.parts { + 0 => format!("fetching {}{from} -- {}", self.model, self.file), + parts => format!( + "fetching {}{from} -- {} ({} of {parts})", + self.model, self.file, self.part + ), + } + } +} + +/// Why a fetch did not come back with what it was asked for. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum Failed { + /// Neither a published name nor a manifest that exists. + UnknownModel(String), + /// The `stop` it was given said to give up. + Stopped, + /// The download failed. + Fetch(String), +} + +impl Failed { + fn from_hub(error: hub::Error) -> Failed { + match hub::stopped(&error) { + true => Failed::Stopped, + false => Failed::Fetch(error.to_string()), + } + } +} + +/// The name a screen shows for `asked`: the file name of a path, since the whole of one does not +/// fit, and the name itself otherwise. +pub fn shown_name(asked: &str) -> String { + match asked.rsplit_once('/') { + Some((_, file)) if !file.is_empty() => file.to_string(), + _ => asked.to_string(), + } +} + +/// Fetches the published model `asked` names, if it is not here already, and hands back its +/// manifest; a manifest's own path is handed back as it is. On the caller's thread. +pub fn fetch( + asked: &str, + on_progress: &mut dyn FnMut(&Fetch), + stop: &dyn Fn() -> bool, +) -> Result { + if hub::full_name(asked).is_none() && !Path::new(asked).is_file() { + return Err(Failed::UnknownModel(format!( + "there is no model called \"{asked}\": neither a published name nor a manifest that \ + exists" + ))); + } + + let mut fetch = Fetch::of(shown_name(asked)); + on_progress(&fetch); + hub::resolve_reporting( + asked, + &mut |progress| { + fetch.heard(progress); + on_progress(&fetch); + }, + stop, + ) + .map_err(Failed::from_hub) +} + +/// Fetches the manifest of every published model not yet here, and none of their packages: see +/// [`hub::fetch_manifests`]. Each is reported as a part of the whole, under the model name +/// "model metadata". +pub fn fetch_manifests( + on_progress: &mut dyn FnMut(&Fetch), + stop: &dyn Fn() -> bool, +) -> Result<(), Failed> { + let mut fetch = Fetch::of("model metadata"); + hub::fetch_manifests( + &mut |progress| { + fetch.heard(progress); + on_progress(&fetch); + }, + stop, + ) + .map_err(Failed::from_hub) +} + +/// The published models, aliases only, as JSON: what `waifu_modelmanager_catalog_json` hands out. +pub fn catalog_json() -> serde_json::Value { + let entry = |listed: hub::Listed, kind: &str| { + serde_json::json!({ + "name": listed.name, + "full_name": listed.full_name, + "kind": kind, + "cached": listed.cached, + // Whether its manifest is here, with or without its packages: what it suggests can + // be read from it. See waifu_modelmanager_fetch_manifests_async. + "manifest_here": hub::local_manifest(listed.name).is_some(), + "bytes_on_disk": listed.bytes, + // Where its packages are, once any of them is: the model's own directory, which + // removing it deletes. + "directory": hub::local_directory(listed.name).map(|directory| directory.to_string_lossy().into_owned()), + "explicit": listed.explicit, + }) + }; + let models: Vec = hub::listed() + .into_iter() + .map(|listed| entry(listed, "image")) + .chain(hub::listed_voices().into_iter().map(|listed| entry(listed, "speech"))) + .chain(hub::listed_conversions().into_iter().map(|listed| entry(listed, "conversion"))) + .collect(); + serde_json::json!({ "models": models }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_fetch_of_an_unknown_length_has_no_fraction() { + // Which is what a server that did not say how long the file is leaves, and it is also + // what the moment before the first byte looks like. + let mut fetch = Fetch { + file: "sdxl-base.waifupkg".to_string(), + done: 512, + part: 1, + parts: 3, + ..Fetch::of("sdxl:base") + }; + assert_eq!(fetch.fraction(), None); + + fetch.total = Some(1024); + assert_eq!(fetch.fraction(), Some(0.5)); + + let words = fetch.words(); + assert!(words.contains("sdxl:base"), "{words}"); + assert!(words.contains("1 of 3"), "{words}"); + } + + #[test] + fn a_name_that_is_neither_published_nor_a_file_is_refused_before_anything_is_fetched() { + let mut heard = 0; + let fetched = fetch("no-such:model", &mut |_| heard += 1, &|| false); + assert!(matches!(fetched, Err(Failed::UnknownModel(_))), "{fetched:?}"); + assert_eq!(heard, 0); + } + + #[test] + fn what_is_published_comes_back_as_json() { + let catalog = catalog_json(); + let models = catalog["models"].as_array().unwrap(); + assert!(models.iter().any(|model| model["name"] == "sdxl:base" && model["kind"] == "image")); + assert!(models.iter().any(|model| model["kind"] == "speech")); + } + + #[test] + fn a_path_shows_as_its_file_name() { + assert_eq!(shown_name("/models/sdxl-base.yaml"), "sdxl-base.yaml"); + assert_eq!(shown_name("sdxl:base"), "sdxl:base"); + } +} diff --git a/waifu/src/hub.rs b/waifu-modelmanager/src/hub.rs similarity index 96% rename from waifu/src/hub.rs rename to waifu-modelmanager/src/hub.rs index 1dade60..ea22809 100644 --- a/waifu/src/hub.rs +++ b/waifu-modelmanager/src/hub.rs @@ -32,9 +32,9 @@ //! a list of files to go and get. The packages are not written down here as well, because a table //! that repeated them could come to disagree with the model. //! -//! This lives in the command line tool rather than in the library: a name is a convenience for -//! someone typing at a terminal, and a program embedding the library has its own ideas about -//! where its files live and when it is allowed to touch the network. +//! This lives in the model manager rather than in the engine: a program embedding the engine may +//! have its own ideas about where its files live and when it is allowed to touch the network, and +//! one that only manages models has no use for the engine at all. use std::env; use std::fmt; @@ -51,17 +51,17 @@ use hf_hub::progress::{DownloadEvent, Progress as Reported, ProgressEvent, Progr use hf_hub::HFClientSync; use crate::config; -use crate::Manifest; +use crate::yaml; -type Error = Box; +pub type Error = Box; /// The suffixes a model's files carry, which is what tells a path from a mistyped name. /// /// A manifest is what `-m` is given; the weights are named here as well because typing one of /// those is an easy enough mistake to make, and it should be answered as the file it is rather /// than as a name nobody published. -const MANIFEST_SUFFIX: &str = Manifest::SUFFIX; -const WEIGHTS_SUFFIX: &str = Manifest::WEIGHTS_SUFFIX; +const MANIFEST_SUFFIX: &str = ".yaml"; +const WEIGHTS_SUFFIX: &str = ".safetensors"; /// Where the cache goes when the environment does not say. const CACHE_ENV: &str = "WAIFU_CACHE"; @@ -653,11 +653,9 @@ pub fn model_bytes(model: &str) -> Option { Some(_) => cached_manifest(model)?, None => PathBuf::from(model), }; - let manifest = Manifest::open(manifest).ok()?; - let mut bytes = 0; - for name in manifest.files() { - bytes += fs::metadata(manifest.file(&name).ok()?).ok()?.len(); + for (_, path) in named_files(&manifest).ok()? { + bytes += fs::metadata(path).ok()?.len(); } Some(bytes) } @@ -869,7 +867,7 @@ fn reads_as_a_name(model: &str) -> bool { /// /// A path is taken as it is written. A known name is fetched into the cache, and what comes back /// is the model's manifest -- the file that says what it is and names its packages, which is what -/// [`Manifest::open`] expects to be handed. +/// a manifest's path is expected to be. /// /// There is no variant that reports nowhere. A fetch is minutes long, and the one caller there /// is has a bar to draw it on; somewhere for it to say so is not optional. @@ -947,20 +945,62 @@ fn fetch(published: &Published, watching: &mut Watching) -> Result Result, Error> { - let manifest = Manifest::open(manifest)?; - let files = manifest.files(); + Ok(named_files(manifest)?.into_iter().map(|(name, _)| name).collect()) +} - for name in &files { - manifest.file(name)?; +/// Every file the manifest at `manifest` names, and where it is: the `weights` list, then each +/// tokenizer's `tokenizer.json` once. What the engine reads out of the same two blocks is +/// `waifu::Manifest::files`; the rest of a manifest is the engine's alone. +/// +/// The names come out of a file fetched over the network and are about to be joined to a +/// directory and written to, so each is checked before any of it is used: one component of a +/// path, a neighbour of the manifest and nothing else -- not `..`, not a directory, not an +/// absolute path to somewhere else on the machine. +fn named_files(manifest: &Path) -> Result, Error> { + let said = |error: &dyn fmt::Display| format!("{}: {error}", manifest.display()); + let text = fs::read_to_string(manifest).map_err(|error| said(&error))?; + let mut top = yaml::parse(&text) + .and_then(|node| node.into_map("the manifest")) + .map_err(|error| said(&error))?; + + let mut names = Vec::new(); + if let Some(weights) = top.remove("weights") { + for item in weights.into_seq("weights")? { + names.push(item.into_scalar("weights")?); + } } - Ok(files) + if names.is_empty() { + return Err(said(&"this manifest names no weights").into()); + } + if let Some(tokenizers) = top.remove("tokenizers") { + for (name, file) in tokenizers.into_map("tokenizers")? { + let file = file.into_scalar(&name)?; + // Two tokenizers of one model may name one file, and a file is fetched once. + if !names.contains(&file) { + names.push(file); + } + } + } + + let beside = manifest.parent().unwrap_or(Path::new(".")); + names + .into_iter() + .map(|name| { + let neighbour = !name.is_empty() + && name != "." + && name != ".." + && !name.contains(['/', '\\']) + && Path::new(&name).components().count() == 1; + match neighbour { + true => Ok((name.clone(), beside.join(&name))), + false => Err(said(&format!( + "this model names {name:?}, and a manifest may only name a neighbour of its own" + )) + .into()), + } + }) + .collect() } /// Which of the three places that can name the model directory named it. @@ -1580,7 +1620,7 @@ mod tests { } /// Ignored because it asks the real Hugging Face for a real file. Run it with - /// `cargo test --features cli -- --ignored` after touching the fetch. + /// `cargo test -p waifu-modelmanager -- --ignored` after touching the fetch. #[test] #[ignore] fn a_file_fetched_from_hugging_face_lands_where_it_was_asked_for() { diff --git a/waifu/src/bin/waifu.rs b/waifu-modelmanager/src/lib.rs similarity index 51% rename from waifu/src/bin/waifu.rs rename to waifu-modelmanager/src/lib.rs index f16f3df..0b51f9b 100644 --- a/waifu/src/bin/waifu.rs +++ b/waifu-modelmanager/src/lib.rs @@ -17,11 +17,38 @@ // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -//! The `waifu` command line tool. +//! The model manager: which models are published, fetching them, and where they are kept. //! -//! A shim so that the implementation can live in the library, where its unit tests run with the -//! rest of the crate. See [`waifu::cli`]. +//! libwaifu_modelmanager. None of it runs a model, so none of it links the tensor library, and the +//! engine, libwaifu, does not link this: an app -- or the command line -- asks this for a model's +//! manifest and hands the engine its path. What anything outside this crate calls is the C API in +//! [`c_api`]; the Rust items are public for this crate's own tests and examples. -fn main() -> std::process::ExitCode { - waifu::cli::run() +pub mod c_api; +pub mod config; +pub mod defaults; +pub mod fetch; +pub mod hub; +pub mod log; +pub mod suggested; +pub mod yaml; + +use std::fmt; + +/// A manifest, or the YAML it is written in, that does not hold what the format calls for. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct FormatError(String); + +impl FormatError { + pub fn new(message: impl Into) -> FormatError { + FormatError(message.into()) + } } + +impl fmt::Display for FormatError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +impl std::error::Error for FormatError {} diff --git a/waifu-modelmanager/src/log.rs b/waifu-modelmanager/src/log.rs new file mode 100644 index 0000000..0f66d6d --- /dev/null +++ b/waifu-modelmanager/src/log.rs @@ -0,0 +1,92 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! Where the log lines go: the console, unless something says otherwise. +//! +//! An app has no console its user would ever read, so it hands over a sink, and from then on every +//! line -- a download that is taking another route, say -- reaches it instead. `waifu::log` is +//! built over this one and sends the tensor library's lines here too, so that in one image there +//! is one sink and one level. + +use std::sync::atomic::{AtomicI32, Ordering}; +use std::sync::RwLock; + +/// How much a line matters. In the tensor library's order, and with its numbers, so that one +/// level set here is the same level there. +#[derive(Clone, Copy, Debug, Eq, PartialEq, PartialOrd, Ord)] +#[repr(i32)] +pub enum Level { + Debug = 0, + Info = 1, + Warning = 2, + Error = 3, + /// The last line: the process aborts once the sink returns. + Fatal = 4, +} + +impl Level { + pub fn name(self) -> &'static str { + match self { + Level::Debug => "DEBUG", + Level::Info => "INFO", + Level::Warning => "WARNING", + Level::Error => "ERROR", + Level::Fatal => "FATAL", + } + } +} + +/// What receives a line: its level, where it was written -- a module's name, or a file and line +/// of the tensor library's -- and the line alone. +/// +/// Called on whichever thread wrote the line, so it has to be safe to call from any of them, and +/// it must not set another sink from inside itself, which would wait on the lock it is called +/// under. +pub type Sink = Box; + +static SINK: RwLock> = RwLock::new(None); +static LEVEL: AtomicI32 = AtomicI32::new(Level::Info as i32); + +/// Sends every line to `sink` from now on, or back to the console for `None`. +pub fn set_sink(sink: Option) { + *SINK.write().unwrap_or_else(|held| held.into_inner()) = sink; +} + +/// Writes no line below `level`. Info unless set. +pub fn set_level(level: Level) { + LEVEL.store(level as i32, Ordering::Relaxed); +} + +/// Writes one line. +pub fn write(level: Level, source: &str, message: &str) { + if (level as i32) < LEVEL.load(Ordering::Relaxed) { + return; + } + hand_over(level, source, message); +} + +/// The line to the sink where there is one, and the console where there is not, without asking +/// the level: for a line whose level was checked already, where it was written. +pub fn hand_over(level: Level, source: &str, message: &str) { + let sink = SINK.read().unwrap_or_else(|held| held.into_inner()); + match sink.as_ref() { + Some(sink) => sink(level, source, message), + None => eprintln!("{} {source}] {message}", level.name()), + } +} diff --git a/waifu/src/suggested.rs b/waifu-modelmanager/src/suggested.rs similarity index 87% rename from waifu/src/suggested.rs rename to waifu-modelmanager/src/suggested.rs index 0546bf5..a47c5bd 100644 --- a/waifu/src/suggested.rs +++ b/waifu-modelmanager/src/suggested.rs @@ -17,7 +17,7 @@ // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -//! What a model suggests being asked for: the `suggested:` block of a [`Manifest`]. +//! What a model suggests being asked for: the `suggested:` block of its manifest. //! //! The rest of the manifest says what the model *is* -- how many layers, how wide, which packages //! the weights are in -- and every key of it has to be there or the model cannot be built. This is @@ -40,8 +40,8 @@ //! //! Every field is optional and is [`None`] when the model does not say, because none of it is //! needed to load a model -- and because "this model has no opinion" is a real answer, distinct -//! from any particular number. What fills the gaps is [`GenerationDefaults`], which is what this -//! build believed before a model could say anything at all. +//! from any particular number. What fills the gaps is the kind of model's own defaults: see +//! [`crate::defaults`]. //! //! A manifest with no `suggested:` block reads as [`Suggestions::default`] rather than as an //! error, and a key this version does not know is skipped rather than refused, so a newer model @@ -49,7 +49,7 @@ use std::collections::BTreeMap; -use crate::generation::GenerationDefaults; +use crate::defaults::Picture; use crate::yaml::Node; /// One size a model says it draws well at, in pixels. @@ -137,6 +137,20 @@ impl Suggestions { } } + /// The `suggested:` block of the manifest at `path`, or nothing suggested where there is none + /// or the file does not read. + pub fn read(path: &std::path::Path) -> Suggestions { + let block = std::fs::read_to_string(path) + .ok() + .and_then(|text| crate::yaml::parse(&text).ok()) + .and_then(|node| node.into_map("the manifest").ok()) + .and_then(|mut top| top.remove("suggested")); + match block { + Some(Node::Map(fields)) => Suggestions::from_fields(&fields), + _ => Suggestions::default(), + } + } + /// The size to start at, which is the first the model lists. pub fn size(&self) -> Option { self.sizes.first().copied() @@ -147,14 +161,14 @@ impl Suggestions { /// The two are asked in that order on purpose: `defaults` is what this build believes about /// the kind of model, and a model that says nothing gets exactly that, which is what every /// model got before there was anywhere to say it. - pub fn over(&self, defaults: GenerationDefaults) -> GenerationDefaults { + pub fn over(&self, defaults: Picture) -> Picture { let (width, height) = self.size().unwrap_or((defaults.width, defaults.height)); - GenerationDefaults { + Picture { width, height, - num_steps: self.steps.unwrap_or(defaults.num_steps), - guidance_scale: self.guidance.unwrap_or(defaults.guidance_scale), + steps: self.steps.unwrap_or(defaults.steps), + guidance: self.guidance.unwrap_or(defaults.guidance), } } } @@ -333,13 +347,31 @@ mod tests { assert_eq!(suggested.prompt.as_deref(), Some("1girl")); } + #[test] + fn a_prompt_and_a_list_are_read_the_way_yaml_reads_them() { + let prompt = |line: &str| Suggestions::from_fields(&fields(&format!("suggested:\n {line}\n"))).prompt; + + // A space and then a `#` begins a comment; a `#` with no space before it is part of the + // value, and a prompt that really wants the first is quoted, which the exporter does. + assert_eq!(prompt("prompt: a plain value # and a comment").as_deref(), Some("a plain value")); + assert_eq!(prompt("prompt: tag#1 keeps its hash").as_deref(), Some("tag#1 keeps its hash")); + assert_eq!(prompt(r#"prompt: "tag #1, quoted""#).as_deref(), Some("tag #1, quoted")); + assert_eq!(prompt("prompt: 'it''s here'").as_deref(), Some("it's here")); + assert_eq!(prompt("prompt: 猫, 🐱").as_deref(), Some("猫, 🐱")); + + let block = Suggestions::from_fields(&fields("suggested:\n sizes:\n - [1024, 1024]\n - [832, 1216]\n")).sizes; + let inline = Suggestions::from_fields(&fields("suggested:\n sizes: [[1024, 1024], [832, 1216]]\n")).sizes; + assert_eq!(block, vec![(1024, 1024), (832, 1216)]); + assert_eq!(block, inline); + } + #[test] fn what_a_model_says_wins_and_what_it_does_not_say_is_left_alone() { - let defaults = GenerationDefaults { + let defaults = Picture { width: 1024, height: 1024, - num_steps: 30, - guidance_scale: 5.0, + steps: 30, + guidance: 5.0, }; // A model with nothing to say gets exactly what this build believed. @@ -352,8 +384,8 @@ mod tests { ..Suggestions::default() }; let applied = turbo.over(defaults); - assert_eq!(applied.num_steps, 8); - assert_eq!(applied.guidance_scale, 1.0); + assert_eq!(applied.steps, 8); + assert_eq!(applied.guidance, 1.0); assert_eq!(applied.width, 1024); assert_eq!(applied.height, 1024); @@ -364,6 +396,6 @@ mod tests { }; let applied = portrait.over(defaults); assert_eq!((applied.width, applied.height), (832, 1216)); - assert_eq!(applied.num_steps, 30); + assert_eq!(applied.steps, 30); } } diff --git a/waifu-modelmanager/src/yaml.rs b/waifu-modelmanager/src/yaml.rs new file mode 100644 index 0000000..4c64239 --- /dev/null +++ b/waifu-modelmanager/src/yaml.rs @@ -0,0 +1,429 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! Enough YAML to read a file this crate also writes the exporter for. +//! +//! Block mappings, block sequences, flow sequences and scalars, indented with spaces. That is +//! what `tools/package_common.py` emits and what a manifest is; everything +//! else in the language -- anchors, flow mappings, tags, multiple documents, block scalars -- is +//! refused rather than half-supported. The one thing worse than a format this cannot read is one +//! it reads differently from everybody else. +//! +//! Every scalar stays the text it was written as. `steps: 8` is the string `"8"` here and becomes +//! a number where it is used, because what a value means is the thing reading it's to say: the +//! same `320,640,1280` is three numbers to the U-Net and a string to anything else. + +use std::collections::BTreeMap; + +use crate::FormatError; + +type Result = std::result::Result; + +/// One value, at whatever depth. +#[derive(Clone, Debug, PartialEq)] +pub enum Node { + Scalar(String), + Map(BTreeMap), + Seq(Vec), +} + +impl Node { + pub fn into_map(self, what: &str) -> Result> { + match self { + Node::Map(map) => Ok(map), + _ => Err(FormatError::new(format!("{what} is not a block of keys"))), + } + } + + pub fn into_seq(self, what: &str) -> Result> { + match self { + Node::Seq(items) => Ok(items), + _ => Err(FormatError::new(format!("{what} is not a list"))), + } + } + + pub fn into_scalar(self, what: &str) -> Result { + match self { + Node::Scalar(text) => Ok(text), + _ => Err(FormatError::new(format!("{what} is not a single value"))), + } + } + + /// This value as text, if it is one. + pub fn as_str(&self) -> Option<&str> { + match self { + Node::Scalar(text) => Some(text), + _ => None, + } + } + + /// The items of this value, if it is a list. + pub fn as_seq(&self) -> Option<&[Node]> { + match self { + Node::Seq(items) => Some(items), + _ => None, + } + } +} + +/// One line that is neither blank nor a comment: how far it is indented, what is on it, and which +/// line of the file it was, for the complaints. +struct Row<'a> { + indent: usize, + text: &'a str, + number: usize, +} + +/// The whole document, which is one block at indent zero. +pub fn parse(text: &str) -> Result { + let mut rows = Vec::new(); + for (index, line) in text.lines().enumerate() { + let number = index + 1; + if line.trim().is_empty() || line.trim_start().starts_with('#') { + continue; + } + // A tab is not indentation in YAML, and a file that used one would nest differently here + // than it does anywhere else. Better to say so than to guess a width for it. + let indent = line.len() - line.trim_start_matches(' ').len(); + if line[indent..].starts_with('\t') { + return Err(FormatError::new(format!( + "line {number}: indented with a tab, and indentation here is spaces" + ))); + } + rows.push(Row { + indent, + text: line[indent..].trim_end(), + number, + }); + } + + if rows.is_empty() { + return Err(FormatError::new("there is nothing in this file")); + } + if rows[0].indent != 0 { + return Err(FormatError::new(format!( + "line {}: the file starts indented", + rows[0].number + ))); + } + + let mut at = 0; + let node = block(&rows, &mut at, 0)?; + if at != rows.len() { + return Err(FormatError::new(format!( + "line {}: this is indented less than the block it is in", + rows[at].number + ))); + } + Ok(node) +} + +/// Every row at `indent`, as one mapping or one sequence. +/// +/// Which of the two it is, the first row says: `- ` begins a sequence and anything else a +/// mapping. Rows indented further belong to the value above them and are read by the call this +/// one makes; a row indented less ends the block and is the caller's. +fn block(rows: &[Row], at: &mut usize, indent: usize) -> Result { + match is_item(rows[*at].text) { + true => sequence(rows, at, indent), + false => mapping(rows, at, indent), + } +} + +/// Whether a row is a `- ` list item rather than a key. +fn is_item(text: &str) -> bool { + text == "-" || text.starts_with("- ") +} + +fn mapping(rows: &[Row], at: &mut usize, indent: usize) -> Result { + let mut map: BTreeMap = BTreeMap::new(); + + while *at < rows.len() && rows[*at].indent >= indent { + let row = &rows[*at]; + if row.indent > indent { + return Err(FormatError::new(format!( + "line {}: this is indented further than the key above it", + row.number + ))); + } + if is_item(row.text) { + return Err(FormatError::new(format!( + "line {}: a list item where a key was expected", + row.number + ))); + } + + let (key, rest) = split_key(row.text, row.number)?; + let number = row.number; + *at += 1; + + let value = if rest.is_empty() { + // `key:` with nothing after it: what it holds is the block below, which has to be + // indented further than the key that introduced it. + if *at < rows.len() && rows[*at].indent > indent { + let deeper = rows[*at].indent; + block(rows, at, deeper)? + } else { + return Err(FormatError::new(format!( + "line {number}: {key:?} is given nothing" + ))); + } + } else { + value(rest, number)? + }; + + if map.insert(key.clone(), value).is_some() { + return Err(FormatError::new(format!( + "line {number}: {key:?} is written twice" + ))); + } + } + + Ok(Node::Map(map)) +} + +fn sequence(rows: &[Row], at: &mut usize, indent: usize) -> Result { + let mut items = Vec::new(); + + while *at < rows.len() && rows[*at].indent >= indent { + let row = &rows[*at]; + if row.indent > indent { + return Err(FormatError::new(format!( + "line {}: this is indented further than the item above it", + row.number + ))); + } + if row.text == "-" { + return Err(FormatError::new(format!( + "line {}: a list item with nothing in it", + row.number + ))); + } + let Some(rest) = row.text.strip_prefix("- ") else { + return Err(FormatError::new(format!( + "line {}: a key where a list item was expected", + row.number + ))); + }; + + let number = row.number; + *at += 1; + items.push(value(rest, number)?); + } + + Ok(Node::Seq(items)) +} + +/// A `key:` and whatever followed it on the line. +/// +/// The colon is the first one outside quotes and outside brackets, so a quoted key holding one is +/// still a key. It has to be followed by a space or by nothing: `a:b` is the scalar `a:b` in +/// YAML, and reading it as a key here would be reading this file differently from every other +/// reader. +fn split_key(text: &str, number: usize) -> Result<(String, &str)> { + let mut quote: Option = None; + let mut depth = 0usize; + + for (index, byte) in text.bytes().enumerate() { + match quote { + Some(open) if byte == open => quote = None, + Some(_) => {} + None if byte == b'"' || byte == b'\'' => quote = Some(byte), + None if byte == b'[' => depth += 1, + None if byte == b']' => depth = depth.saturating_sub(1), + None if byte == b':' && depth == 0 => { + let after = &text[index + 1..]; + if !after.is_empty() && !after.starts_with(' ') { + continue; + } + let key = plain_or_quoted(text[..index].trim_end(), number)?; + if key.is_empty() { + return Err(FormatError::new(format!("line {number}: a key with no name"))); + } + return Ok((key, after.trim_start())); + } + None => {} + } + } + Err(FormatError::new(format!( + "line {number}: {text:?} is neither a key nor a list item" + ))) +} + +/// One value as written after a `key:` or a `- `: a flow sequence, or a scalar. +fn value(text: &str, number: usize) -> Result { + let text = text.trim(); + if text.starts_with('[') { + let (node, rest) = flow(text, number)?; + // Trailing text after a closed bracket is a line that means something else somewhere + // else, and guessing which would be the wrong kind of helpful. + let rest = strip_comment(rest); + if !rest.trim().is_empty() { + return Err(FormatError::new(format!( + "line {number}: {rest:?} follows a list and is not part of it" + ))); + } + return Ok(node); + } + Ok(Node::Scalar(plain_or_quoted(text, number)?)) +} + +/// A `[a, b, [c, d]]`, from its opening bracket, and whatever is left of the line after it. +fn flow(text: &str, number: usize) -> Result<(Node, &str)> { + let mut rest = text + .strip_prefix('[') + .expect("a flow sequence starts with its bracket") + .trim_start(); + let mut items = Vec::new(); + + if let Some(after) = rest.strip_prefix(']') { + return Ok((Node::Seq(items), after)); + } + + loop { + let (item, after) = if rest.starts_with('[') { + flow(rest, number)? + } else { + let end = flow_item_ends(rest, number)?; + ( + Node::Scalar(plain_or_quoted(rest[..end].trim(), number)?), + &rest[end..], + ) + }; + items.push(item); + + rest = after.trim_start(); + if let Some(after) = rest.strip_prefix(',') { + rest = after.trim_start(); + // `[1, 2, ]` -- a comma before the close is a trailing comma, not an empty item. + if let Some(after) = rest.strip_prefix(']') { + return Ok((Node::Seq(items), after)); + } + continue; + } + if let Some(after) = rest.strip_prefix(']') { + return Ok((Node::Seq(items), after)); + } + return Err(FormatError::new(format!( + "line {number}: a list that is not closed" + ))); + } +} + +/// Where one scalar item of a flow sequence ends: at the comma or the bracket that follows it, +/// whichever comes first outside quotes. +fn flow_item_ends(text: &str, number: usize) -> Result { + let mut quote: Option = None; + for (index, byte) in text.bytes().enumerate() { + match quote { + Some(open) if byte == open => quote = None, + Some(_) => {} + None if byte == b'"' || byte == b'\'' => quote = Some(byte), + None if byte == b',' || byte == b']' => return Ok(index), + None => {} + } + } + Err(FormatError::new(format!( + "line {number}: a list that is not closed" + ))) +} + +/// A scalar with its quotes taken off, or a plain one with its comment taken off. +fn plain_or_quoted(text: &str, number: usize) -> Result { + let text = text.trim(); + + if let Some(body) = quoted(text, b'\'') { + // A doubled quote is the only escape a single quoted scalar has, so a lone one is a value + // that ended earlier than it looks like it did. + if body.replace("''", "").contains('\'') { + return Err(FormatError::new(format!( + "line {number}: a quote inside a single quoted value has to be doubled" + ))); + } + return Ok(body.replace("''", "'")); + } + + if let Some(body) = quoted(text, b'"') { + return escapes(body, number); + } + + if text.starts_with('"') || text.starts_with('\'') { + return Err(FormatError::new(format!( + "line {number}: {text:?} opens a quote it does not close" + ))); + } + + Ok(strip_comment(text).trim_end().to_string()) +} + +/// What is left of a plain scalar once a trailing comment is taken off. +/// +/// A `#` only begins a comment when a space separates it from what came before, which is the rule +/// everywhere else: `tag#1` is a value holding a hash, and `tag #1` is the value `tag` with a +/// comment after it. A value that really does want a space and then a hash has to be quoted, and +/// the exporter quotes it -- see `_quote` in `tools/package_common.py`. +fn strip_comment(text: &str) -> &str { + for (index, _) in text.match_indices('#') { + if index == 0 || text.as_bytes()[index - 1] == b' ' { + return &text[..index]; + } + } + text +} + +/// The inside of `text`, if the whole of it is wrapped in `mark`. +fn quoted(text: &str, mark: u8) -> Option<&str> { + let bytes = text.as_bytes(); + match bytes.len() >= 2 && bytes[0] == mark && bytes[bytes.len() - 1] == mark { + true => Some(&text[1..text.len() - 1]), + false => None, + } +} + +/// What a double quoted scalar means, which is the only place a backslash does anything. +fn escapes(body: &str, number: usize) -> Result { + let mut text = String::new(); + let mut chars = body.chars(); + while let Some(c) = chars.next() { + if c != '\\' { + text.push(c); + continue; + } + match chars.next() { + Some('n') => text.push('\n'), + Some('t') => text.push('\t'), + Some('r') => text.push('\r'), + Some('"') => text.push('"'), + Some('\'') => text.push('\''), + Some('\\') => text.push('\\'), + Some('/') => text.push('/'), + Some(other) => { + return Err(FormatError::new(format!( + "line {number}: \\{other} is not an escape this reads" + ))) + } + None => { + return Err(FormatError::new(format!( + "line {number}: a value ends in a backslash" + ))) + } + } + } + Ok(text) +} diff --git a/waifu-tools/Cargo.toml b/waifu-tools/Cargo.toml new file mode 100644 index 0000000..f01f1f6 --- /dev/null +++ b/waifu-tools/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "waifu-tools" +version = "0.1.0" +edition = "2021" +license = "MIT" +description = "Tools for working on libwaifu: checking an export, timing a model, fetching manifests" + +# Programs for the people working on this repository, not for anybody using the libraries: unlike +# the command line, they call the two crates' Rust APIs directly -- a model's graph, its +# parameters, the allocator's high-water mark -- which no C API says. Each is one file under +# src/bin, run with `cargo run --release -p waifu-tools --bin -- ...`. + +[dependencies] +waifu = { path = "../waifu" } +waifu-modelmanager = { path = "../waifu-modelmanager" } diff --git a/waifu-ffi/build.rs b/waifu-tools/build.rs similarity index 59% rename from waifu-ffi/build.rs rename to waifu-tools/build.rs index 0fcb213..1f10e3e 100644 --- a/waifu-ffi/build.rs +++ b/waifu-tools/build.rs @@ -17,12 +17,11 @@ // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -//! Links the library with what only a final link can be given, and writes waifu.h. +//! Links what only a final link can be given. //! -//! The waifu crate's build script names flint's static libraries, and those reach this link on +//! The waifu crate's build script names flint's static libraries, and those reach these links on //! their own. Its link arguments do not: cargo keeps a `rustc-link-arg` to the package that said -//! it, so whatever flint's link flags carry as one is said again here, for this package's own -//! links. (The Metal kernels need none: they are bytes inside flint itself.) +//! it, so whatever flint's link flags carry as one is said again here, for these binaries. use std::path::PathBuf; @@ -45,34 +44,4 @@ fn main() { for line in flags.lines().filter(|line| line.starts_with("cargo:rustc-link-arg=")) { println!("{line}"); } - - // Found next to whatever loads it -- an app's Frameworks directory, a test's own -- rather than - // at the path cargo happened to write it to. - if std::env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("macos") { - println!("cargo:rustc-cdylib-link-arg=-Wl,-install_name,@rpath/libwaifu.dylib"); - } - - write_header(&manifest); -} - -/// waifu.h, from the `#[repr(C)]` types and `extern "C"` functions in src/lib.rs, with their doc -/// comments. Written beside the source, where a change to the ABI shows up in a diff. -fn write_header(manifest: &PathBuf) { - println!("cargo:rerun-if-changed=src/lib.rs"); - println!("cargo:rerun-if-changed=cbindgen.toml"); - - let config = cbindgen::Config::from_file(manifest.join("cbindgen.toml")) - .expect("waifu-ffi/cbindgen.toml"); - match cbindgen::Builder::new() - .with_crate(manifest) - .with_config(config) - .generate() - { - Ok(bindings) => { - bindings.write_to_file(manifest.join("include/waifu.h")); - } - // Said rather than failed on: a header that cannot be regenerated leaves the last one, - // and the library under it still builds. - Err(error) => println!("cargo:warning=waifu.h was not regenerated: {error}"), - } } diff --git a/waifu/examples/cosyvoice3_bench.rs b/waifu-tools/src/bin/cosyvoice3_bench.rs similarity index 97% rename from waifu/examples/cosyvoice3_bench.rs rename to waifu-tools/src/bin/cosyvoice3_bench.rs index 77c5345..c48d053 100644 --- a/waifu/examples/cosyvoice3_bench.rs +++ b/waifu-tools/src/bin/cosyvoice3_bench.rs @@ -20,7 +20,7 @@ //! Time CosyVoice3's stages, through nothing but its public API. //! //! ```text -//! cargo run --release --example cosyvoice3_bench -- models/cosyvoice3.yaml voice.wav [runs] +//! cargo run --release -p waifu-tools --bin cosyvoice3_bench -- models/cosyvoice3.yaml voice.wav [runs] //! ``` //! //! Listening to the recording, and then for each of three texts: the language model (from the diff --git a/waifu/examples/fetch_manifests.rs b/waifu-tools/src/bin/fetch_manifests.rs similarity index 89% rename from waifu/examples/fetch_manifests.rs rename to waifu-tools/src/bin/fetch_manifests.rs index 0531dcd..538187e 100644 --- a/waifu/examples/fetch_manifests.rs +++ b/waifu-tools/src/bin/fetch_manifests.rs @@ -22,8 +22,7 @@ //! be shown before anything is fetched (see `hub::set_bundled_manifests`). //! //! ```bash -//! cargo run --release --manifest-path waifu/Cargo.toml --features hub \ -//! --example fetch_manifests -- macos/build/manifests +//! cargo run --release -p waifu-tools --bin fetch_manifests -- waifu-app/swiftui/build/manifests //! ``` use std::path::PathBuf; @@ -40,9 +39,9 @@ fn main() { // The cache is where fetch_manifests writes; pointed here, it writes nowhere else. std::env::set_var("WAIFU_CACHE", &directory); - let fetched = waifu::hub::fetch_manifests( + let fetched = waifu_modelmanager::hub::fetch_manifests( &mut |progress| { - if let waifu::hub::Progress::Fetched { file, part, parts, .. } = progress { + if let waifu_modelmanager::hub::Progress::Fetched { file, part, parts, .. } = progress { eprintln!("{part}/{parts} {file}"); } }, diff --git a/waifu/examples/generate.rs b/waifu-tools/src/bin/generate.rs similarity index 83% rename from waifu/examples/generate.rs rename to waifu-tools/src/bin/generate.rs index a370f84..cb010d6 100644 --- a/waifu/examples/generate.rs +++ b/waifu-tools/src/bin/generate.rs @@ -20,7 +20,7 @@ //! Draws one picture and writes it out as a PPM, which is a header and the pixels. //! //! ```text -//! cargo run --release --example generate -- sdxl.yaml "an astronaut riding a horse on mars" +//! cargo run --release -p waifu-tools --bin generate -- sdxl.yaml "an astronaut riding a horse on mars" //! ``` //! //! A third argument names the device -- cpu, cuda or metal. Left out, it takes the first @@ -38,10 +38,8 @@ use std::io::Write; -use waifu::{ - to_rgb8, Anima, Device, GenerationDefaults, GenerationOptions, Krea2, Manifest, Residency, - Sdxl, -}; +use waifu::{to_rgb8, Anima, Device, Krea2, Manifest, Residency, Sdxl}; +use waifu_tools::settings; fn main() -> Result<(), waifu::Error> { let mut arguments: Vec = std::env::args().skip(1).collect(); @@ -85,22 +83,9 @@ fn main() -> Result<(), waifu::Error> { let kind = manifest.section("model")?.get_str("type")?.to_string(); eprintln!("drawing a {kind} on {device:?}, weights {residency:?}"); - // What the kind of model wants, and then what this one says over it. A model with no - // `suggested:` block gets exactly its kind's defaults, which is what every model got before - // there was anywhere to say otherwise. - let defaults = match kind.as_str() { - Anima::MODEL_TYPE => Anima::DEFAULTS, - Krea2::MODEL_TYPE => Krea2::DEFAULTS, - _ => GenerationDefaults::default(), - }; - let suggested = manifest.suggested(); - let options = GenerationOptions { - // The author's own, where the card gave one. A model that suggests nothing to steer away - // from leaves this empty, which is what it always was here. - negative_prompt: suggested.avoid.clone().unwrap_or_default(), - seed: Some(7), - ..suggested.over(defaults).options() - }; + // What the kind of model wants, and then what this one says over it, as the model manager + // works it out for a screen. A model with no `suggested:` block gets exactly its kind's. + let options = settings(&model_path)?; eprintln!( "{} by {}, {} steps, guidance {}", options.width, options.height, options.num_steps, options.guidance_scale diff --git a/waifu/examples/indextts_bench.rs b/waifu-tools/src/bin/indextts_bench.rs similarity index 98% rename from waifu/examples/indextts_bench.rs rename to waifu-tools/src/bin/indextts_bench.rs index 9dc48e0..0785b0f 100644 --- a/waifu/examples/indextts_bench.rs +++ b/waifu-tools/src/bin/indextts_bench.rs @@ -20,7 +20,7 @@ //! Time IndexTTS-2.5's stages, through nothing but its public API. //! //! ```text -//! cargo run --release --example indextts_bench -- models/indextts25.yaml voice.wav [runs] [out] +//! cargo run --release -p waifu-tools --bin indextts_bench -- models/indextts25.yaml voice.wav [runs] [out] //! ``` //! //! Listening to the recording, and then for each of three texts: the GPT (from the start of a diff --git a/waifu/examples/inspect.rs b/waifu-tools/src/bin/inspect.rs similarity index 94% rename from waifu/examples/inspect.rs rename to waifu-tools/src/bin/inspect.rs index 8aeb7a0..7fc83e7 100644 --- a/waifu/examples/inspect.rs +++ b/waifu-tools/src/bin/inspect.rs @@ -21,7 +21,7 @@ //! hold. //! //! ```text -//! cargo run --release --example inspect -- models/sdxl-base.yaml +//! cargo run --release -p waifu-tools --bin inspect -- models/sdxl-base.yaml //! ``` //! //! With `-n` it stops after the manifest and reads no weights, which is the quick way to see what @@ -44,7 +44,7 @@ fn main() -> Result<(), waifu::Error> { let model_type = manifest.section("model")?.get_str("type")?.to_string(); println!("kind: {model_type}"); - let suggested = manifest.suggested(); + let suggested = waifu_modelmanager::suggested::Suggestions::read(manifest.path()); if let Some(prompt) = &suggested.prompt { println!("suggested prompt: {prompt:?}"); } diff --git a/waifu/examples/load.rs b/waifu-tools/src/bin/load.rs similarity index 97% rename from waifu/examples/load.rs rename to waifu-tools/src/bin/load.rs index d08fd86..8b31151 100644 --- a/waifu/examples/load.rs +++ b/waifu-tools/src/bin/load.rs @@ -20,7 +20,7 @@ //! Builds a whole model from its manifest and encodes a prompt with it, without drawing anything. //! //! ```text -//! cargo run --release --example load -- models/sdxl-base.yaml [cpu|cuda] +//! cargo run --release -p waifu-tools --bin load -- models/sdxl-base.yaml [cpu|cuda] //! ``` //! //! What `inspect` does not: the weights are not merely read but assembled into the model the diff --git a/waifu/examples/memprofile.rs b/waifu-tools/src/bin/memprofile.rs similarity index 92% rename from waifu/examples/memprofile.rs rename to waifu-tools/src/bin/memprofile.rs index 36da4f1..28e9094 100644 --- a/waifu/examples/memprofile.rs +++ b/waifu-tools/src/bin/memprofile.rs @@ -17,7 +17,8 @@ use std::time::Instant; use waifu::flint::MemorySnapshot; -use waifu::{Device, GenerationOptions, Krea2, Manifest, Residency}; +use waifu::{Device, Krea2, Manifest, Residency}; +use waifu_tools::settings; const PROMPT: &str = "a calico cat asleep on a stack of books, warm afternoon light"; @@ -66,12 +67,7 @@ fn main() -> Result<(), waifu::Error> { build_time ); - let suggested = manifest.suggested(); - let options = GenerationOptions { - negative_prompt: suggested.avoid.clone().unwrap_or_default(), - seed: Some(7), - ..suggested.over(Krea2::DEFAULTS).options() - }; + let options = settings(model_path)?; println!( "drawing {} by {}, {} steps, guidance {}", options.width, options.height, options.num_steps, options.guidance_scale diff --git a/waifu/examples/parameters.rs b/waifu-tools/src/bin/parameters.rs similarity index 98% rename from waifu/examples/parameters.rs rename to waifu-tools/src/bin/parameters.rs index b44ca1b..f52c1cc 100644 --- a/waifu/examples/parameters.rs +++ b/waifu-tools/src/bin/parameters.rs @@ -20,7 +20,7 @@ //! Every tensor a model's graph loads, with the shape it expects, as one name per line. //! //! ```text -//! cargo run --release --example parameters -- s2mel +//! cargo run --release -p waifu-tools --bin parameters -- s2mel //! ``` //! //! This is the list an exporter has to satisfy, and it is worth having as a program rather than diff --git a/waifu/examples/speak.rs b/waifu-tools/src/bin/speak.rs similarity index 93% rename from waifu/examples/speak.rs rename to waifu-tools/src/bin/speak.rs index f2bba30..60b12fc 100644 --- a/waifu/examples/speak.rs +++ b/waifu-tools/src/bin/speak.rs @@ -20,9 +20,9 @@ //! Say a sentence in the voice of a recording, with IndexTTS-2.5 or Fun-CosyVoice3. //! //! ```text -//! cargo run --release --example speak -- models/indextts25.yaml voice.wav "Hello there." out.wav -//! cargo run --release --example speak -- models/cosyvoice3.yaml voice.wav "你好。" out.wav 7 "它的文字稿。" -//! cargo run --release --example speak -- models/cosyvoice3.yaml voice.wav "你好。" out.wav 7 --style=请非常开心地说一句话。 +//! cargo run --release -p waifu-tools --bin speak -- models/indextts25.yaml voice.wav "Hello there." out.wav +//! cargo run --release -p waifu-tools --bin speak -- models/cosyvoice3.yaml voice.wav "你好。" out.wav 7 "它的文字稿。" +//! cargo run --release -p waifu-tools --bin speak -- models/cosyvoice3.yaml voice.wav "你好。" out.wav 7 --style=请非常开心地说一句话。 //! ``` //! //! The whole pipeline from the command line, with nothing of the web page in the way: what each diff --git a/waifu-tools/src/lib.rs b/waifu-tools/src/lib.rs new file mode 100644 index 0000000..2adf75a --- /dev/null +++ b/waifu-tools/src/lib.rs @@ -0,0 +1,41 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! What the tools in src/bin share. + +use waifu::GenerationOptions; + +/// The request a tool draws with for the model whose manifest is at `path`: what a screen would +/// start its boxes at -- the kind of model's defaults with the manifest's `suggested:` block over +/// them, as libwaifu_modelmanager works it out -- at a fixed seed, so two runs draw one picture. +pub fn settings(path: &str) -> Result { + let said = waifu_modelmanager::defaults::defaults_json(path) + .ok_or_else(|| waifu::Error::Model(format!("there is no manifest at {path}")))?; + let number = |key: &str| said[key].as_f64().unwrap_or_default(); + Ok(GenerationOptions { + width: number("width") as i32, + height: number("height") as i32, + num_steps: number("steps") as i32, + guidance_scale: number("guidance") as f32, + // The author's own, where the card gave one; nothing where it did not. + negative_prompt: said["avoid"].as_str().unwrap_or_default().to_string(), + seed: Some(7), + ..GenerationOptions::default() + }) +} diff --git a/waifu/Cargo.toml b/waifu/Cargo.toml index 3cc7b11..b0e11f1 100644 --- a/waifu/Cargo.toml +++ b/waifu/Cargo.toml @@ -8,93 +8,30 @@ description = "Your local AI waifu: diffusion image generation on top of the fli # to a single owner in the dependency graph. links = "flint" +# A dynamic library, libwaifu.dylib (.so, waifu.dll), whose C API is all anything outside this +# crate uses: see docs/ffi.md. The Metal kernels are a section of the image flint is in, and a +# dylib carries them wherever it goes. The rlib is for this crate's own tests, for the command +# line, which links it statically and calls its C API, and for waifu-tools. [lib] name = "waifu" path = "src/lib.rs" - -[[bin]] -name = "waifu" -path = "src/bin/waifu.rs" -required-features = ["cli"] - -[features] -# The command line tool. Off by default: what it pulls in -- a terminal stack, an HTTP client, an -# image encoder -- is the binary's to need, and nothing the library does asks for any of it. -cli = [ - "engine", - "image/png", - "image/jpeg", - "dep:ratatui", -] -# Fetching published models, and where they are kept: the catalogue, the two hubs, config.toml. -hub = ["dep:hf-hub", "dep:toml", "dep:ureq"] -# What a program that is not the command line needs to run models for somebody: the engine, what -# it says about models and the machine as JSON, and the resizing a picture to start from takes. -# What the C API in waifu-ffi is built over. -engine = ["hub", "dep:image", "dep:serde_json"] +crate-type = ["cdylib", "rlib"] [dependencies] -# What the command line tool needs. -# -# The terminal the whole of it happens in: which task, which model -- fetched there, under a bar -# -- and which device, and then the screen a picture is drawn from. Trimmed to the crossterm -# backend: the default set also builds the termwiz and termina ones, which is two more terminal -# stacks than a program that runs in one terminal needs. -ratatui = { version = "0.29", default-features = false, features = [ - "crossterm", - "underline-color", -], optional = true } - # What the engine says about models and the machine, for the C API to hand on: a parser this -# crate would own for no reason, since it is already in the tree, under hf-hub. +# crate would own for no reason, since it is already in the tree, under tokenizers. # # `Value` rather than derived types. What crosses is a handful of shapes, each built once, and a # derive would buy nothing at the cost of a second description of every one of them. -serde_json = { version = "1", optional = true } - -# The settings file, `config.toml`, which today says one thing: where fetched models go. It is a -# file for people to open in an editor, so it is in the format that lets them write a comment and -# a Windows path without escaping every backslash -- and reading that properly is a parser, which -# is not something to write again here. -toml = { version = "1", optional = true } - -# Fetching from ModelScope, which hf-hub does not speak, and probing whether the machine can -# reach the outside at all. Rust has no HTTP in its standard library, so this is the one thing the -# tool cannot do on its own. Held to the rustls backend: the native-tls default would link the -# system OpenSSL, and a dynamic library on the far side of the build is exactly what this crate -# does not want. -ureq = { version = "3", default-features = false, features = [ - "rustls", - "gzip", -], optional = true } - -# Fetching from Hugging Face, which is where the models are published and where all but one of -# the fetches go. Its own client rather than ureq against a URL built here: it is what tracks how -# that hub actually answers -- redirects to a CDN, xet-backed files that are not a plain GET, the -# headers that say a file moved -- and none of that is worth rediscovering here one bug report at -# a time. -# -# It is the largest thing in this file by a distance, reqwest and tokio and hf-xet behind it, and -# that is the price of the above. Held to rustls for the same reason ureq is, and to the blocking -# feature because nothing else in this crate is async: an executor exists to run one download at -# a time and no further. -# -# The other fetch is ModelScope, which this does not speak. That one stays on ureq, below. -hf-hub = { version = "1", default-features = false, features = [ - "blocking", - "rustls-tls", -], optional = true } +serde_json = "1" -# Reading the picture an image to image run starts from, and writing the one it ends with. The -# default feature set is a dozen decoders -- avif, exr, tiff, webp and the rest -- behind a good -# deal of C. Held to the two formats anyone actually hands a diffusion model, which are also the -# two written in pure Rust. +# Resizing the picture an image to image run starts from to the size it is drawn at. No codecs: +# a picture crosses the C API as rows of RGB, and decoding or writing a file is the caller's. # # Cargo.lock holds 0.25.9 rather than the newest 0.25: 0.25.10 asks for a rustc a minor version # ahead of the one this is built with. Nothing here needs what it added, so the lock is the place # that says so rather than a bound in the range. -# No codecs by default: the engine only resizes, and the command line turns on the two it reads. -image = { version = "0.25", default-features = false, optional = true } +image = { version = "0.25", default-features = false } # Reading the weights, which are safetensors: the format every other runtime already writes, so a # weight can be looked at with torch, with numpy, or with the viewer on the hub. It replaced two @@ -122,7 +59,5 @@ tokenizers = { version = "0.23", default-features = false, features = [ "fancy-regex", ] } -# Fetches every published model's manifest into a directory, for an app to be built with. -[[example]] -name = "fetch_manifests" -required-features = ["hub"] +[build-dependencies] +cbindgen = { version = "0.29", default-features = false } diff --git a/waifu/build.rs b/waifu/build.rs index 75ac99a..585ea27 100644 --- a/waifu/build.rs +++ b/waifu/build.rs @@ -1,6 +1,6 @@ -// Points the linker at the static library that the CMake build produces. CMake is what drives -// this crate's build -- the top-level CMakeLists.txt's `waifu-cli` target runs `cargo build` after -// the `flint` archive is fresh, so a plain `cmake --build build` builds both. This script never +// Points the linker at the static library that the CMake build produces, and writes waifu.h. CMake +// is what drives this crate's build -- the top-level CMakeLists.txt's `waifu-engine` target runs +// `cargo build` after the `flint` archive is fresh, so a plain `cmake --build build` builds both. This script never // invokes CMake itself; it only reads the link flags CMake already wrote out. That keeps a lone // `cargo build` working too (useful when iterating on Rust code only), as long as the native // archive is already up to date. Override the location with LIBWAIFU_LIB_DIR when building @@ -45,46 +45,36 @@ fn main() { println!("{line}"); } - stamp_revision(); -} + // Found next to whatever loads it -- an app's Frameworks directory, the command line's own -- + // rather than at the path cargo happened to write it to. + if std::env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("macos") { + println!("cargo:rustc-cdylib-link-arg=-Wl,-install_name,@rpath/libwaifu.dylib"); + } -/// Put the commit this was built from into the binary, for the screen to show. -/// -/// A screenshot of a run is worth little without it: two builds of the same version can differ by -/// anything, and the first thing to ask about a picture that came out wrong is which code made it. -/// Missing on purpose is a failure: a copy of the source with no git around it still builds, and -/// says "unknown" where the hash would be rather than refusing. -fn stamp_revision() { let manifest = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); - let git = manifest.join("../.git"); + write_header(&manifest); +} + +/// waifu.h, from the `extern "C"` functions and `#[repr(C)]` types in src/c_api.rs and +/// src/c_api/support.rs, with their doc comments. Written beside the source, where a change to the +/// ABI shows up in a diff. +fn write_header(manifest: &Path) { + println!("cargo:rerun-if-changed=src/c_api.rs"); + println!("cargo:rerun-if-changed=src/c_api/support.rs"); + println!("cargo:rerun-if-changed=cbindgen.toml"); - // Rebuilt when the checkout moves. HEAD covers a checkout, and the file HEAD points at covers - // a commit made on the branch already checked out. - if git.is_dir() { - println!("cargo:rerun-if-changed={}", git.join("HEAD").display()); - if let Some(reference) = head_reference(&git) { - println!("cargo:rerun-if-changed={}", git.join(reference).display()); + let config = cbindgen::Config::from_file(manifest.join("cbindgen.toml")).expect("waifu/cbindgen.toml"); + match cbindgen::Builder::new() + .with_config(config) + .with_src(manifest.join("src/c_api.rs")) + .with_src(manifest.join("src/c_api/support.rs")) + .generate() + { + Ok(bindings) => { + bindings.write_to_file(manifest.join("include/waifu.h")); } + // Said rather than failed on: a header that cannot be regenerated leaves the last one, + // and the library under it still builds. + Err(error) => println!("cargo:warning=waifu.h was not regenerated: {error}"), } - - let revision = std::process::Command::new("git") - .args(["rev-parse", "--short", "HEAD"]) - .current_dir(&manifest) - .output() - .ok() - .filter(|out| out.status.success()) - .and_then(|out| String::from_utf8(out.stdout).ok()) - .map(|text| text.trim().to_string()) - .filter(|text| !text.is_empty()) - .unwrap_or_else(|| "unknown".to_string()); - - println!("cargo:rustc-env=WAIFU_REVISION={revision}"); -} - -/// The ref file HEAD names, relative to the git directory, when it names one at all. A detached -/// HEAD holds a commit instead, and then there is no second file to watch. -fn head_reference(git: &Path) -> Option { - let head = std::fs::read_to_string(git.join("HEAD")).ok()?; - let reference = head.trim().strip_prefix("ref: ")?; - Some(reference.to_string()) } diff --git a/waifu-ffi/cbindgen.toml b/waifu/cbindgen.toml similarity index 68% rename from waifu-ffi/cbindgen.toml rename to waifu/cbindgen.toml index 1544a46..a605386 100644 --- a/waifu-ffi/cbindgen.toml +++ b/waifu/cbindgen.toml @@ -1,4 +1,7 @@ -# How waifu.h is written from src/lib.rs. See docs/ffi.md for the contract it keeps. +# How waifu.h is written from src/c_api.rs. See docs/ffi.md for the contract it keeps. +# +# Everything the header declares is this library's own: it includes nothing but the C library's, +# and shares no type with libwaifu_modelmanager's. language = "C" include_guard = "LIBWAIFU_WAIFU_H" @@ -10,17 +13,17 @@ sys_includes = ["stdbool.h", "stddef.h", "stdint.h"] no_includes = true # size_t for a length, as C callers expect and Swift imports as Int. usize_is_size_t = true -autogen_warning = "/* Written by cbindgen from waifu-ffi/src/lib.rs. Change that, not this. */" +autogen_warning = "/* Written by cbindgen from waifu/src/c_api.rs. Change that, not this. */" header = """/* The MIT License (MIT). Copyright (c) 2026 Xiaoyang Chen. See LICENSE. */""" [defines] [enum] -# The variants are named in src/lib.rs as C names them: WAIFU_OK, not WaifuStatusCode_Ok. rename_variants = "None" prefix_with_name = false [export] +# Every type of the engine's, whether or not a function names it directly. include = [ "WaifuStatusCode", "WaifuStatus", @@ -29,6 +32,7 @@ include = [ "WaifuStage", "WaifuProgressEvent", "WaifuCompletionEvent", + "WaifuDevice", "WaifuImage", "WaifuImageInput", "WaifuDrawRequest", @@ -37,5 +41,4 @@ include = [ "WaifuSpeakRequest", "WaifuVoiceConversionRequest", "WaifuAudioCompletionEvent", - "WaifuDevice", ] diff --git a/waifu-ffi/include/waifu.h b/waifu/include/waifu.h similarity index 75% rename from waifu-ffi/include/waifu.h rename to waifu/include/waifu.h index d0744aa..937dea4 100644 --- a/waifu-ffi/include/waifu.h +++ b/waifu/include/waifu.h @@ -3,7 +3,7 @@ #ifndef LIBWAIFU_WAIFU_H #define LIBWAIFU_WAIFU_H -/* Written by cbindgen from waifu-ffi/src/lib.rs. Change that, not this. */ +/* Written by cbindgen from waifu/src/c_api.rs. Change that, not this. */ #include #include @@ -11,7 +11,7 @@ // What a caller compiled against this header expects. Bumped when anything in it changes in a // way an older caller would misread. -#define WAIFU_ABI_VERSION 1 +#define WAIFU_ABI_VERSION 2 typedef enum { WAIFU_LOG_DEBUG = 0, @@ -22,29 +22,22 @@ typedef enum { WAIFU_LOG_FATAL = 4, } WaifuLogLevel; -// What a call or a job came to. typedef enum { - WAIFU_OK = 0, - WAIFU_CANCELLED = 1, - // A NULL, a string that is not UTF-8, a size out of range, a request that names an image - // it does not give. - WAIFU_ERR_INVALID_ARGUMENT = 2, - // Neither a published name nor a manifest that exists. - WAIFU_ERR_UNKNOWN_MODEL = 3, - // The download failed. - WAIFU_ERR_FETCH = 4, - // The model could not be read, or cannot do what was asked. - WAIFU_ERR_MODEL = 5, - // A job with no model loaded, or a model of the wrong kind. - WAIFU_ERR_NOT_LOADED = 6, - // A panic in the library: a bug. - WAIFU_ERR_INTERNAL = 7, -} WaifuStatusCode; + WAIFU_DEVICE_AUTO = 0, + WAIFU_DEVICE_CPU = 1, + WAIFU_DEVICE_METAL = 2, + WAIFU_DEVICE_CUDA = 3, + WAIFU_DEVICE_CUDA_CPU_OFFLOAD = 4, + WAIFU_DEVICE_VULKAN = 5, + // Metal with the weights left in the model's files, each read off the disk as it is used: + // for a model larger than the Mac's memory, and slower. Last, so the numbers above keep + // their meaning. + WAIFU_DEVICE_METAL_DISK_OFFLOAD = 6, +} WaifuDevice; typedef enum { - // file, done/total in bytes, part/parts. - WAIFU_STAGE_FETCHING = 0, - // Reading weights onto the device: no fraction. + // Reading weights onto the device: no fraction. 0 was fetching, which is + // libwaifu_modelmanager's now; the others keep their numbers. WAIFU_STAGE_READING = 1, // The prompt, or the picture img2img starts from. WAIFU_STAGE_ENCODING = 2, @@ -59,25 +52,28 @@ typedef enum { WAIFU_STAGE_SOUNDING = 7, } WaifuStage; +// What a call or a job came to. typedef enum { - WAIFU_DEVICE_AUTO = 0, - WAIFU_DEVICE_CPU = 1, - WAIFU_DEVICE_METAL = 2, - WAIFU_DEVICE_CUDA = 3, - WAIFU_DEVICE_CUDA_CPU_OFFLOAD = 4, - WAIFU_DEVICE_VULKAN = 5, - // Metal with the weights left in the model's files, each read off the disk as it is used: - // for a model larger than the Mac's memory, and slower. Last, so the numbers above keep - // their meaning. - WAIFU_DEVICE_METAL_DISK_OFFLOAD = 6, -} WaifuDevice; + WAIFU_OK = 0, + WAIFU_CANCELLED = 1, + // A NULL, a string that is not UTF-8, a size out of range, a request that names an image + // it does not give. + WAIFU_ERR_INVALID_ARGUMENT = 2, + // Neither a published name nor a manifest that exists. + WAIFU_ERR_UNKNOWN_MODEL = 3, + // The download failed. + WAIFU_ERR_FETCH = 4, + // The model could not be read, or cannot do what was asked. + WAIFU_ERR_MODEL = 5, + // A job with no model loaded, or a model of the wrong kind. + WAIFU_ERR_NOT_LOADED = 6, + // A panic in the library: a bug. + WAIFU_ERR_INTERNAL = 7, +} WaifuStatusCode; // One thread that holds a model and runs jobs on it, one at a time, in the order asked. typedef struct WaifuEngine WaifuEngine; -// A download running on a thread of its own. -typedef struct WaifuModelFetch WaifuModelFetch; - typedef struct { WaifuLogLevel level; // Where it was written: "interface.cc:84" for flint, a module's name for Rust. @@ -90,6 +86,11 @@ typedef struct { // call from any thread, and must not call into the library. typedef void (*WaifuLogCallback)(void *user_data, const WaifuLogEvent *event); +// 0 is never a job: an _async call returns it when it refuses at the call, with the reason in +// waifu_last_error(), and then calls none of its callbacks. Any other id gets exactly one +// on_complete. +typedef uint64_t WaifuJob; + typedef struct { WaifuStage stage; // 0..1 of the whole job, or -1 where nothing can say. @@ -97,10 +98,10 @@ typedef struct { // What the stage counts; 0, 0 where it counts nothing. uint64_t done; uint64_t total; - // FETCHING only: which file of how many. + // Reserved: 0, 0 and NULL. They said which file of a download, which is + // libwaifu_modelmanager's now, and are kept so the event keeps its layout. uint32_t part; uint32_t parts; - // FETCHING only, else NULL. const char *file; // "step 3 of 8", for a status line. const char *words; @@ -124,11 +125,6 @@ typedef struct { typedef void (*WaifuCompleteCallback)(void *user_data, const WaifuCompletionEvent *event); -// 0 is never a job: an _async call returns it when it refuses at the call, with the reason in -// waifu_last_error(), and then calls none of its callbacks. Any other id gets exactly one -// on_complete. -typedef uint64_t WaifuJob; - // A picture as rows of RGB. typedef struct { uint32_t width; @@ -232,7 +228,7 @@ extern "C" { // The ABI this library was built with: WAIFU_ABI_VERSION. uint32_t waifu_abi_version(void); -// This thread's last message; valid until this thread's next call. +// This thread's last message from this library; valid until this thread's next call into it. const char *waifu_last_error(void); // Frees a string a synchronous call returned. NULL is fine. @@ -250,50 +246,17 @@ void waifu_set_log_level(WaifuLogLevel level); // with waifu_string_free. NULL on failure. char *waifu_machine_json(void); -// The published models, aliases only, as JSON. Free with waifu_string_free. NULL on failure. -char *waifu_modelmanager_catalog_json(void); - -// What a model is, without reading it, as JSON: a published name or a manifest path. Free with -// waifu_string_free. NULL, with waifu_last_error, for a name that is neither. -char *waifu_modelmanager_describe_json(const char *model); - -// Where downloads go. Free with waifu_string_free. NULL on failure. -char *waifu_modelmanager_directory(void); - -// Downloads go to `path` from now on, saved to config.toml; NULL or "" for the default. -WaifuStatusCode waifu_modelmanager_set_directory(const char *path); - -// Reads the manifests of published models out of `directory` where the download directory has -// none: an app built with every model's manifest (the `fetch_manifests` example fetches them) -// can say what each model suggests before anything is downloaded. Laid out as the download -// directory is, and only read. NULL or "" for none. Not saved. -WaifuStatusCode waifu_modelmanager_set_bundled_manifests(const char *directory); - -// Deletes a downloaded package. -WaifuStatusCode waifu_modelmanager_remove(const char *name); - -// Starts fetching a published model on a thread of its own and returns at once. Fetching one -// already here finishes at once. Where it went is waifu_modelmanager_describe_json's to say. -// NULL, with waifu_last_error and no callback ever, where it cannot start. -WaifuModelFetch *waifu_modelmanager_fetch_async(const char *model, - void *user_data, - WaifuProgressCallback on_progress, - WaifuCompleteCallback on_complete); - -// Starts fetching the manifest of every published model that is not here yet -- a couple of -// kilobytes each, and none of their packages -- on a thread of its own, and returns at once. -// What a model suggests is in its manifest, so a screen that has called this can fill in a -// model's settings the moment it is chosen. Each manifest is a part of the whole in the progress; -// with every one already here it finishes at once. Cancelled and freed as a model's fetch is. -// NULL, with waifu_last_error and no callback ever, where it cannot start. -WaifuModelFetch *waifu_modelmanager_fetch_manifests_async(void *user_data, - WaifuProgressCallback on_progress, - WaifuCompleteCallback on_complete); - -void waifu_modelmanager_fetch_cancel(WaifuModelFetch *fetch); +// Every device, as JSON: `{"devices": [{"name", "available", "free"}]}`, named as the -device flag +// names them -- cpu, cuda, cuda_cpu_offload, metal, metal_disk_offload, vulkan -- with `free` the +// bytes free on the card for CUDA and Vulkan where they can be used, and null otherwise. Asking +// CUDA builds a context on the card, a few hundred megabytes of it. Free with waifu_string_free. +// NULL on failure. +char *waifu_devices_json(void); -// After on_complete; cancels first if it has not had it, and waits for the thread to end. -void waifu_modelmanager_fetch_free(WaifuModelFetch *fetch); +// What a model is, without reading it, as JSON: the path of its manifest, whose weights need not +// be beside it, or "tones". Free with waifu_string_free. NULL, with waifu_last_error, for +// anything else. +char *waifu_describe_json(const char *model); // Starts the engine's thread. NULL, with waifu_last_error, if the device is not available. WaifuEngine *waifu_engine_new(WaifuDevice device); @@ -307,8 +270,9 @@ void waifu_engine_cancel(WaifuEngine *engine, WaifuJob job); void waifu_engine_cancel_all(WaifuEngine *engine); -// Fetches if needed, then reads onto the device, dropping the model before. What was loaded is -// waifu_modelmanager_describe_json's to say. +// Reads the model whose manifest is at `model` onto the device, dropping the one before; "tones" +// is the stand-in voice. A published model is fetched, and its manifest found, with +// libwaifu_modelmanager first. What was loaded is waifu_describe_json's to say. WaifuJob waifu_engine_load_async(WaifuEngine *engine, const char *model, void *user_data, diff --git a/waifu-ffi/src/lib.rs b/waifu/src/c_api.rs similarity index 55% rename from waifu-ffi/src/lib.rs rename to waifu/src/c_api.rs index 74e6448..71e4185 100644 --- a/waifu-ffi/src/lib.rs +++ b/waifu/src/c_api.rs @@ -17,7 +17,9 @@ // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -//! The C API over libwaifu: the engine, the model manager and the log, as `extern "C"`. +//! The C API: the engine, what a model is, the machine and the log, as `extern "C"`. Everything +//! outside this crate goes through these. A model is the path of its manifest; listing, fetching +//! and removing published models is libwaifu_modelmanager's, which shares nothing with this. //! //! Nothing with any logic in it lives here. This converts -- C strings to `&str`, structs to the //! engine's requests, what comes back to borrowed C views, panics to status codes -- and keeps @@ -26,63 +28,28 @@ #![allow(non_camel_case_types)] #![allow(clippy::missing_safety_doc)] -use std::cell::RefCell; -use std::ffi::{c_char, c_void, CStr, CString}; -use std::panic::{self, AssertUnwindSafe}; +use std::ffi::{c_char, c_void, CString}; use std::ptr; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Condvar, Mutex}; -use std::thread::{self, JoinHandle}; -use waifu::engine::{self, DrawRequest, Ended, Engine, Failure, SpeakRequest, Stage, VoiceConversionRequest}; -use waifu::runtime::DeviceOption; -use waifu::wav::Sound; +mod support; -/// What a caller compiled against this header expects. Bumped when anything in it changes in a -/// way an older caller would misread. -pub const WAIFU_ABI_VERSION: u32 = 1; - -// -- status ---------------------------------------------------------------------------------- +pub use support::{ + WaifuCompleteCallback, WaifuCompletionEvent, WaifuLogCallback, WaifuLogEvent, WaifuLogLevel, + WaifuProgressCallback, WaifuProgressEvent, WaifuStage, WaifuStatus, WaifuStatusCode, +}; +use support::{caught, required, set_last_error, status, string_in, string_out, Gate, UserData}; -/// What a call or a job came to. -#[repr(C)] -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum WaifuStatusCode { - WAIFU_OK = 0, - WAIFU_CANCELLED = 1, - /// A NULL, a string that is not UTF-8, a size out of range, a request that names an image - /// it does not give. - WAIFU_ERR_INVALID_ARGUMENT = 2, - /// Neither a published name nor a manifest that exists. - WAIFU_ERR_UNKNOWN_MODEL = 3, - /// The download failed. - WAIFU_ERR_FETCH = 4, - /// The model could not be read, or cannot do what was asked. - WAIFU_ERR_MODEL = 5, - /// A job with no model loaded, or a model of the wrong kind. - WAIFU_ERR_NOT_LOADED = 6, - /// A panic in the library: a bug. - WAIFU_ERR_INTERNAL = 7, -} +use crate::engine::{self, DrawRequest, Ended, Engine, Failure, SpeakRequest, Stage, VoiceConversionRequest}; +use crate::runtime::DeviceOption; +use crate::wav::Sound; use WaifuStatusCode::*; -/// How a job ended, as on_complete is told it. -#[repr(C)] -pub struct WaifuStatus { - pub code: WaifuStatusCode, - /// Why, when code is not WAIFU_OK; NULL when it is. - pub message: *const c_char, -} - -thread_local! { - static LAST_ERROR: RefCell = RefCell::new(CString::default()); -} +/// What a caller compiled against this header expects. Bumped when anything in it changes in a +/// way an older caller would misread. +pub const WAIFU_ABI_VERSION: u32 = 2; -fn set_last_error(message: &str) { - let message = CString::new(message.replace('\0', " ")).unwrap_or_default(); - LAST_ERROR.with(|last| *last.borrow_mut() = message); -} +// -- status ---------------------------------------------------------------------------------- /// The ABI this library was built with: WAIFU_ABI_VERSION. #[no_mangle] @@ -90,130 +57,32 @@ pub extern "C" fn waifu_abi_version() -> u32 { WAIFU_ABI_VERSION } -/// This thread's last message; valid until this thread's next call. +/// This thread's last message from this library; valid until this thread's next call into it. #[no_mangle] pub extern "C" fn waifu_last_error() -> *const c_char { - LAST_ERROR.with(|last| last.borrow().as_ptr()) + support::last_error() } /// Frees a string a synchronous call returned. NULL is fine. #[no_mangle] pub unsafe extern "C" fn waifu_string_free(s: *mut c_char) { - if !s.is_null() { - drop(CString::from_raw(s)); - } -} - -/// Runs a synchronous entry point, turning a panic into WAIFU_ERR_INTERNAL with its message. -fn caught(otherwise: T, body: impl FnOnce() -> T) -> T { - match panic::catch_unwind(AssertUnwindSafe(body)) { - Ok(value) => value, - Err(cause) => { - let said = cause - .downcast_ref::<&str>() - .map(|said| said.to_string()) - .or_else(|| cause.downcast_ref::().cloned()) - .unwrap_or_else(|| "it panicked".to_string()); - set_last_error(&format!("a bug in the library: {said}")); - otherwise - } - } -} - -/// A string handed in, or the reason it cannot be read. NULL is `allow_null`'s to decide. -unsafe fn string_in(pointer: *const c_char, what: &str) -> Result, String> { - if pointer.is_null() { - return Ok(None); - } - CStr::from_ptr(pointer) - .to_str() - .map(|text| Some(text.to_string())) - .map_err(|_| format!("{what} is not UTF-8")) -} - -unsafe fn required(pointer: *const c_char, what: &str) -> Result { - string_in(pointer, what)?.ok_or_else(|| format!("{what} is NULL")) -} - -fn string_out(text: &str) -> *mut c_char { - CString::new(text.replace('\0', " ")).unwrap_or_default().into_raw() + support::free_string(s) } // -- logging --------------------------------------------------------------------------------- -#[repr(C)] -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum WaifuLogLevel { - WAIFU_LOG_DEBUG = 0, - WAIFU_LOG_INFO = 1, - WAIFU_LOG_WARNING = 2, - WAIFU_LOG_ERROR = 3, - /// The last line: the process aborts when the callback returns. - WAIFU_LOG_FATAL = 4, -} - -#[repr(C)] -pub struct WaifuLogEvent { - pub level: WaifuLogLevel, - /// Where it was written: "interface.cc:84" for flint, a module's name for Rust. - pub source: *const c_char, - /// The line alone, with no level or time in front of it. - pub message: *const c_char, -} - -/// Called on whichever thread wrote the line, which can be the caller's own. It must be safe to -/// call from any thread, and must not call into the library. -pub type WaifuLogCallback = Option; - -/// A caller's pointer, handed back as it was given, on whichever thread a callback runs. -#[derive(Clone, Copy)] -struct UserData(*mut c_void); -unsafe impl Send for UserData {} -unsafe impl Sync for UserData {} - /// Every line the library would print -- flint's and the Rust side's -- goes to the callback /// instead, from whichever thread wrote it. NULL puts them back on the console. Set it before /// anything else: flint writes what hardware it found the first time a device is asked about. #[no_mangle] pub extern "C" fn waifu_set_log_callback(user_data: *mut c_void, callback: WaifuLogCallback) { - caught((), || { - let Some(callback) = callback else { - waifu::log::set_sink(None); - return; - }; - let user_data = UserData(user_data); - waifu::log::set_sink(Some(Box::new(move |level, source, message| { - let source = CString::new(source.replace('\0', " ")).unwrap_or_default(); - let message = CString::new(message.replace('\0', " ")).unwrap_or_default(); - let event = WaifuLogEvent { - level: match level { - waifu::log::Level::Debug => WaifuLogLevel::WAIFU_LOG_DEBUG, - waifu::log::Level::Info => WaifuLogLevel::WAIFU_LOG_INFO, - waifu::log::Level::Warning => WaifuLogLevel::WAIFU_LOG_WARNING, - waifu::log::Level::Error => WaifuLogLevel::WAIFU_LOG_ERROR, - waifu::log::Level::Fatal => WaifuLogLevel::WAIFU_LOG_FATAL, - }, - source: source.as_ptr(), - message: message.as_ptr(), - }; - let user_data = user_data; - unsafe { callback(user_data.0, &event) }; - }))); - }) + caught((), || crate::log::set_sink(support::log_sink(user_data, callback))) } /// Lines below the level are not made at all. INFO unless set. #[no_mangle] pub extern "C" fn waifu_set_log_level(level: WaifuLogLevel) { - caught((), || { - waifu::log::set_level(match level { - WaifuLogLevel::WAIFU_LOG_DEBUG => waifu::log::Level::Debug, - WaifuLogLevel::WAIFU_LOG_INFO => waifu::log::Level::Info, - WaifuLogLevel::WAIFU_LOG_WARNING => waifu::log::Level::Warning, - WaifuLogLevel::WAIFU_LOG_ERROR => waifu::log::Level::Error, - WaifuLogLevel::WAIFU_LOG_FATAL => waifu::log::Level::Fatal, - }) - }) + caught((), || crate::log::set_level(support::log_level(level))) } // -- the machine ----------------------------------------------------------------------------- @@ -225,84 +94,20 @@ pub extern "C" fn waifu_machine_json() -> *mut c_char { caught(ptr::null_mut(), || string_out(&engine::machine_json().to_string())) } -// -- progress -------------------------------------------------------------------------------- - -#[repr(C)] -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum WaifuStage { - /// file, done/total in bytes, part/parts. - WAIFU_STAGE_FETCHING = 0, - /// Reading weights onto the device: no fraction. - WAIFU_STAGE_READING = 1, - /// The prompt, or the picture img2img starts from. - WAIFU_STAGE_ENCODING = 2, - /// done/total in steps. - WAIFU_STAGE_DRAWING = 3, - WAIFU_STAGE_DECODING = 4, - /// A conversion reading its recordings. - WAIFU_STAGE_LISTENING = 5, - /// done/total in tokens; total is an estimate and can be passed. - WAIFU_STAGE_SAYING = 6, - /// The vocoder. - WAIFU_STAGE_SOUNDING = 7, -} - -#[repr(C)] -pub struct WaifuProgressEvent { - pub stage: WaifuStage, - /// 0..1 of the whole job, or -1 where nothing can say. - pub fraction: f64, - /// What the stage counts; 0, 0 where it counts nothing. - pub done: u64, - pub total: u64, - /// FETCHING only: which file of how many. - pub part: u32, - pub parts: u32, - /// FETCHING only, else NULL. - pub file: *const c_char, - /// "step 3 of 8", for a status line. - pub words: *const c_char, - /// Since the job started. - pub seconds: f64, -} - -pub type WaifuProgressCallback = Option; - -/// How a job that makes nothing ended: a fetch, a load, an unload. -#[repr(C)] -pub struct WaifuCompletionEvent { - pub status: WaifuStatus, +/// Every device, as JSON: `{"devices": [{"name", "available", "free"}]}`, named as the -device flag +/// names them -- cpu, cuda, cuda_cpu_offload, metal, metal_disk_offload, vulkan -- with `free` the +/// bytes free on the card for CUDA and Vulkan where they can be used, and null otherwise. Asking +/// CUDA builds a context on the card, a few hundred megabytes of it. Free with waifu_string_free. +/// NULL on failure. +#[no_mangle] +pub extern "C" fn waifu_devices_json() -> *mut c_char { + caught(ptr::null_mut(), || string_out(&engine::devices_json().to_string())) } -pub type WaifuCompleteCallback = Option; - -// -- keeping the promises about when a callback can arrive --------------------------------------- - -/// Shut until the `_async` call that made it has returned: a callback that arrives before then -/// waits. The engine's thread can finish a job before the call that queued it is done handing -/// back its id, and a caller whose continuation is set up from that id would miss the end. -#[derive(Clone, Default)] -struct Gate(Arc<(Mutex, Condvar)>); - -impl Gate { - fn open(&self) { - let (open, opened) = &*self.0; - *open.lock().unwrap_or_else(|held| held.into_inner()) = true; - opened.notify_all(); - } - - fn wait(&self) { - let (open, opened) = &*self.0; - let mut is_open = open.lock().unwrap_or_else(|held| held.into_inner()); - while !*is_open { - is_open = opened.wait(is_open).unwrap_or_else(|held| held.into_inner()); - } - } -} +// -- progress -------------------------------------------------------------------------------- fn stage_out(stage: Stage) -> WaifuStage { match stage { - Stage::Fetching => WaifuStage::WAIFU_STAGE_FETCHING, Stage::Reading => WaifuStage::WAIFU_STAGE_READING, Stage::Encoding => WaifuStage::WAIFU_STAGE_ENCODING, Stage::Drawing => WaifuStage::WAIFU_STAGE_DRAWING, @@ -317,25 +122,25 @@ fn stage_out(stage: Stage) -> WaifuStage { fn progress_to(gate: &Gate, user_data: UserData, callback: WaifuProgressCallback) -> engine::OnProgress { let gate = gate.clone(); Box::new(move |progress| { - let Some(callback) = callback else { + if callback.is_none() { return; - }; + } gate.wait(); - let file = progress.file.as_deref().map(|file| CString::new(file.replace('\0', " ")).unwrap_or_default()); - let words = CString::new(progress.words.replace('\0', " ")).unwrap_or_default(); - let event = WaifuProgressEvent { - stage: stage_out(progress.stage), - fraction: progress.fraction.unwrap_or(-1.0), - done: progress.done, - total: progress.total, - part: progress.part, - parts: progress.parts, - file: file.as_ref().map_or(ptr::null(), |file| file.as_ptr()), - words: words.as_ptr(), - seconds: progress.seconds, - }; - let user_data = user_data; - unsafe { callback(user_data.0, &event) }; + support::tell_progress( + callback, + user_data, + &support::Progress { + stage: stage_out(progress.stage), + fraction: progress.fraction, + done: progress.done, + total: progress.total, + part: 0, + parts: 0, + file: None, + words: &progress.words, + seconds: progress.seconds, + }, + ); }) } @@ -348,23 +153,15 @@ fn status_of(ended: &Ended) -> (WaifuStatusCode, Option) { Ended::Failed(failure, message) => ( match failure { Failure::UnknownModel => WAIFU_ERR_UNKNOWN_MODEL, - Failure::Fetch => WAIFU_ERR_FETCH, Failure::Model => WAIFU_ERR_MODEL, Failure::NotLoaded => WAIFU_ERR_NOT_LOADED, Failure::Internal => WAIFU_ERR_INTERNAL, }, - Some(CString::new(message.replace('\0', " ")).unwrap_or_default()), + Some(support::c_string(message)), ), } } -fn status(code: WaifuStatusCode, message: &Option) -> WaifuStatus { - WaifuStatus { - code, - message: message.as_ref().map_or(ptr::null(), |message| message.as_ptr()), - } -} - fn completion_to(gate: &Gate, user_data: UserData, callback: WaifuCompleteCallback) -> engine::OnComplete { let gate = gate.clone(); Box::new(move |ended: Ended| { @@ -381,18 +178,13 @@ fn completion_to(gate: &Gate, user_data: UserData, callback: WaifuCompleteCal }) } -// -- the model manager ----------------------------------------------------------------------- - -/// The published models, aliases only, as JSON. Free with waifu_string_free. NULL on failure. -#[no_mangle] -pub extern "C" fn waifu_modelmanager_catalog_json() -> *mut c_char { - caught(ptr::null_mut(), || string_out(&engine::catalog_json().to_string())) -} +// -- what a model is ------------------------------------------------------------------------- -/// What a model is, without reading it, as JSON: a published name or a manifest path. Free with -/// waifu_string_free. NULL, with waifu_last_error, for a name that is neither. +/// What a model is, without reading it, as JSON: the path of its manifest, whose weights need not +/// be beside it, or "tones". Free with waifu_string_free. NULL, with waifu_last_error, for +/// anything else. #[no_mangle] -pub unsafe extern "C" fn waifu_modelmanager_describe_json(model: *const c_char) -> *mut c_char { +pub unsafe extern "C" fn waifu_describe_json(model: *const c_char) -> *mut c_char { caught(ptr::null_mut(), || { let model = match required(model, "model") { Ok(model) => model, @@ -404,223 +196,145 @@ pub unsafe extern "C" fn waifu_modelmanager_describe_json(model: *const c_char) match engine::describe_json(&model) { Some(described) => string_out(&described.to_string()), None => { - set_last_error(&format!("there is no model called \"{model}\"")); + set_last_error(&crate::describe::no_manifest_at(&model)); ptr::null_mut() } } }) } -/// Where downloads go. Free with waifu_string_free. NULL on failure. -#[no_mangle] -pub extern "C" fn waifu_modelmanager_directory() -> *mut c_char { - caught(ptr::null_mut(), || match waifu::hub::model_directory() { - Ok((directory, _)) => string_out(&directory.to_string_lossy()), - Err(error) => { - set_last_error(&error.to_string()); - ptr::null_mut() - } - }) -} +// -- the engine ------------------------------------------------------------------------------ -/// Downloads go to `path` from now on, saved to config.toml; NULL or "" for the default. -#[no_mangle] -pub unsafe extern "C" fn waifu_modelmanager_set_directory(path: *const c_char) -> WaifuStatusCode { - caught(WAIFU_ERR_INTERNAL, || { - let path = match string_in(path, "path") { - Ok(path) => path.filter(|path| !path.is_empty()), - Err(error) => { - set_last_error(&error); - return WAIFU_ERR_INVALID_ARGUMENT; - } - }; - match waifu::config::set_model_dir(path.as_deref().map(std::path::Path::new)) { - Ok(_) => WAIFU_OK, - Err(error) => { - set_last_error(&error.to_string()); - WAIFU_ERR_INTERNAL - } - } - }) +/// One thread that holds a model and runs jobs on it, one at a time, in the order asked. +pub struct WaifuEngine { + engine: Engine, } -/// Reads the manifests of published models out of `directory` where the download directory has -/// none: an app built with every model's manifest (the `fetch_manifests` example fetches them) -/// can say what each model suggests before anything is downloaded. Laid out as the download -/// directory is, and only read. NULL or "" for none. Not saved. -#[no_mangle] -pub unsafe extern "C" fn waifu_modelmanager_set_bundled_manifests(directory: *const c_char) -> WaifuStatusCode { - caught(WAIFU_ERR_INTERNAL, || { - let directory = match string_in(directory, "directory") { - Ok(directory) => directory.filter(|directory| !directory.is_empty()), - Err(error) => { - set_last_error(&error); - return WAIFU_ERR_INVALID_ARGUMENT; - } - }; - waifu::hub::set_bundled_manifests(directory.map(std::path::PathBuf::from)); - WAIFU_OK - }) -} +/// 0 is never a job: an _async call returns it when it refuses at the call, with the reason in +/// waifu_last_error(), and then calls none of its callbacks. Any other id gets exactly one +/// on_complete. +pub type WaifuJob = u64; -/// Deletes a downloaded package. -#[no_mangle] -pub unsafe extern "C" fn waifu_modelmanager_remove(name: *const c_char) -> WaifuStatusCode { - caught(WAIFU_ERR_INTERNAL, || { - let name = match required(name, "name") { - Ok(name) => name, - Err(error) => { - set_last_error(&error); - return WAIFU_ERR_INVALID_ARGUMENT; - } - }; - if waifu::hub::full_name(&name).is_none() { - set_last_error(&format!("there is no published model called \"{name}\"")); - return WAIFU_ERR_UNKNOWN_MODEL; - } - match waifu::hub::remove(&name) { - Ok(()) => WAIFU_OK, - Err(error) => { - set_last_error(&error.to_string()); - WAIFU_ERR_INTERNAL - } - } - }) +#[repr(C)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum WaifuDevice { + WAIFU_DEVICE_AUTO = 0, + WAIFU_DEVICE_CPU = 1, + WAIFU_DEVICE_METAL = 2, + WAIFU_DEVICE_CUDA = 3, + WAIFU_DEVICE_CUDA_CPU_OFFLOAD = 4, + WAIFU_DEVICE_VULKAN = 5, + /// Metal with the weights left in the model's files, each read off the disk as it is used: + /// for a model larger than the Mac's memory, and slower. Last, so the numbers above keep + /// their meaning. + WAIFU_DEVICE_METAL_DISK_OFFLOAD = 6, } -/// A download running on a thread of its own. -pub struct WaifuModelFetch { - stop: Arc, - thread: Option>, +/// A picture as rows of RGB. +#[repr(C)] +pub struct WaifuImage { + pub width: u32, + pub height: u32, + /// width * height * 3 bytes, rows top to bottom. + pub rgb: *const u8, + pub len: usize, } -/// Starts fetching a published model on a thread of its own and returns at once. Fetching one -/// already here finishes at once. Where it went is waifu_modelmanager_describe_json's to say. -/// NULL, with waifu_last_error and no callback ever, where it cannot start. -#[no_mangle] -pub unsafe extern "C" fn waifu_modelmanager_fetch_async( - model: *const c_char, - user_data: *mut c_void, - on_progress: WaifuProgressCallback, - on_complete: WaifuCompleteCallback, -) -> *mut WaifuModelFetch { - caught(ptr::null_mut(), || { - let model = match required(model, "model") { - Ok(model) => model, - Err(error) => { - set_last_error(&error); - return ptr::null_mut(); - } - }; - spawn_fetch(user_data, on_progress, on_complete, move |on_progress, stop| { - engine::fetch(&model, on_progress, stop) - }) - }) +/// An image a run is given, by name. A name the prompt writes as <|name|> is read as part of +/// what it is told; "start_from" (img2img: see strength) and "control" (a ControlNet's condition: +/// see control_scale) are kept for images that are drawn with, and the prompt cannot name them. +#[repr(C)] +pub struct WaifuImageInput { + /// Letters, digits and _. + pub key: *const c_char, + /// Any size, scaled by the engine. + pub image: *const WaifuImage, } -/// Starts fetching the manifest of every published model that is not here yet -- a couple of -/// kilobytes each, and none of their packages -- on a thread of its own, and returns at once. -/// What a model suggests is in its manifest, so a screen that has called this can fill in a -/// model's settings the moment it is chosen. Each manifest is a part of the whole in the progress; -/// with every one already here it finishes at once. Cancelled and freed as a model's fetch is. -/// NULL, with waifu_last_error and no callback ever, where it cannot start. -#[no_mangle] -pub unsafe extern "C" fn waifu_modelmanager_fetch_manifests_async( - user_data: *mut c_void, - on_progress: WaifuProgressCallback, - on_complete: WaifuCompleteCallback, -) -> *mut WaifuModelFetch { - caught(ptr::null_mut(), || { - spawn_fetch(user_data, on_progress, on_complete, engine::fetch_manifests) - }) +/// Refused at the call, as WAIFU_ERR_INVALID_ARGUMENT: two images with one key; a <|name|> in the +/// prompt with no image by that name, or naming a kept one; an image the prompt does not name and +/// whose key is not a kept one; and any other <|...|>, which the model would read as one of its +/// own markers. +#[repr(C)] +pub struct WaifuDrawRequest { + /// sizeof(WaifuDrawRequest), as the caller was compiled with. + pub struct_size: u32, + pub prompt: *const c_char, + /// Words only; "" or NULL for none; ignored by a model without guidance. + pub negative: *const c_char, + /// NULL, 0 for none. + pub images: *const WaifuImageInput, + pub image_count: usize, + /// 64..2048, multiples of the model's `alignment` in its description (16; 32 for SDXL). + pub width: u32, + pub height: u32, + pub steps: u32, + pub guidance: f32, + /// The caller's: the same seed and settings draw the same image. + pub seed: u64, + /// With "start_from": how far it walks away from it, 0..1. + pub strength: f32, + /// With "control": how hard it holds to it. + pub control_scale: f32, } -/// Runs `work` on a thread of its own, with the callbacks it reports to and a stop it can be -/// cancelled by, and hands back the handle that stop is reached through. -fn spawn_fetch( - user_data: *mut c_void, - on_progress: WaifuProgressCallback, - on_complete: WaifuCompleteCallback, - work: impl FnOnce(&mut engine::OnProgress, &dyn Fn() -> bool) -> Ended + Send + 'static, -) -> *mut WaifuModelFetch { - let gate = Gate::default(); - let user_data = UserData(user_data); - let mut on_progress = progress_to(&gate, user_data, on_progress); - let on_complete = completion_to::(&gate, user_data, on_complete); - let stop = Arc::new(AtomicBool::new(false)); - - let thread = thread::Builder::new().name("waifu-fetch".to_string()).spawn({ - let stop = Arc::clone(&stop); - move || { - let ended = panic::catch_unwind(AssertUnwindSafe(|| { - work(&mut on_progress, &|| stop.load(Ordering::Relaxed)) - })) - .unwrap_or_else(|_| Ended::Failed(Failure::Internal, "a bug in the library: the fetch panicked".to_string())); - on_complete(ended); - } - }); - match thread { - Ok(thread) => { - let fetch = Box::into_raw(Box::new(WaifuModelFetch { - stop, - thread: Some(thread), - })); - gate.open(); - fetch - } - Err(error) => { - set_last_error(&format!("could not start the download's thread: {error}")); - ptr::null_mut() - } - } +#[repr(C)] +pub struct WaifuImageCompletionEvent { + pub status: WaifuStatus, + /// NULL unless status.code is WAIFU_OK. + pub image: *const WaifuImage, } -#[no_mangle] -pub unsafe extern "C" fn waifu_modelmanager_fetch_cancel(fetch: *mut WaifuModelFetch) { - if let Some(fetch) = fetch.as_ref() { - fetch.stop.store(true, Ordering::Relaxed); - } -} +pub type WaifuImageCompleteCallback = + Option; -/// After on_complete; cancels first if it has not had it, and waits for the thread to end. -#[no_mangle] -pub unsafe extern "C" fn waifu_modelmanager_fetch_free(fetch: *mut WaifuModelFetch) { - if fetch.is_null() { - return; - } - let mut fetch = Box::from_raw(fetch); - fetch.stop.store(true, Ordering::Relaxed); - if let Some(thread) = fetch.thread.take() { - let _ = thread.join(); - } +/// The samples, not a file: what goes in -- a recording to sound like, or to convert -- and what +/// comes out. Decoding a file to this and writing this to a file are the caller's. +#[repr(C)] +pub struct WaifuAudio { + /// Mono, -1..1. + pub samples: *const f32, + pub count: usize, + pub rate: u32, } -// -- the engine ------------------------------------------------------------------------------ +#[repr(C)] +pub struct WaifuSpeakRequest { + /// sizeof(WaifuSpeakRequest), as the caller was compiled with. + pub struct_size: u32, + pub text: *const c_char, + /// A recording to sound like, or NULL. + pub like: *const WaifuAudio, + pub speed: f32, + pub temperature: f32, + /// NULL or "" for the voice's own way. + pub style: *const c_char, + pub seed: u64, +} #[repr(C)] -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum WaifuDevice { - WAIFU_DEVICE_AUTO = 0, - WAIFU_DEVICE_CPU = 1, - WAIFU_DEVICE_METAL = 2, - WAIFU_DEVICE_CUDA = 3, - WAIFU_DEVICE_CUDA_CPU_OFFLOAD = 4, - WAIFU_DEVICE_VULKAN = 5, - /// Metal with the weights left in the model's files, each read off the disk as it is used: - /// for a model larger than the Mac's memory, and slower. Last, so the numbers above keep - /// their meaning. - WAIFU_DEVICE_METAL_DISK_OFFLOAD = 6, +pub struct WaifuVoiceConversionRequest { + /// sizeof(WaifuVoiceConversionRequest), as the caller was compiled with. + pub struct_size: u32, + /// What is said. + pub source: *const WaifuAudio, + /// Whose voice to say it in. + pub reference: *const WaifuAudio, + pub steps: u32, + pub convert_style: bool, + pub seed: u64, } -/// One thread that holds a model and runs jobs on it, one at a time, in the order asked. -pub struct WaifuEngine { - engine: Engine, +/// Shared by speak and voice_conversion: both make audio. +#[repr(C)] +pub struct WaifuAudioCompletionEvent { + pub status: WaifuStatus, + /// NULL unless status.code is WAIFU_OK. + pub audio: *const WaifuAudio, } -/// 0 is never a job: an _async call returns it when it refuses at the call, with the reason in -/// waifu_last_error(), and then calls none of its callbacks. Any other id gets exactly one -/// on_complete. -pub type WaifuJob = u64; +pub type WaifuAudioCompleteCallback = + Option; /// Starts the engine's thread. NULL, with waifu_last_error, if the device is not available. #[no_mangle] @@ -687,8 +401,9 @@ fn submitted(gate: &Gate, submit: impl FnOnce() -> Result) -> } } -/// Fetches if needed, then reads onto the device, dropping the model before. What was loaded is -/// waifu_modelmanager_describe_json's to say. +/// Reads the model whose manifest is at `model` onto the device, dropping the one before; "tones" +/// is the stand-in voice. A published model is fetched, and its manifest found, with +/// libwaifu_modelmanager first. What was loaded is waifu_describe_json's to say. #[no_mangle] pub unsafe extern "C" fn waifu_engine_load_async( engine: *mut WaifuEngine, @@ -725,64 +440,6 @@ pub unsafe extern "C" fn waifu_engine_unload_async( // -- images ---------------------------------------------------------------------------------- -/// A picture as rows of RGB. -#[repr(C)] -pub struct WaifuImage { - pub width: u32, - pub height: u32, - /// width * height * 3 bytes, rows top to bottom. - pub rgb: *const u8, - pub len: usize, -} - -/// An image a run is given, by name. A name the prompt writes as <|name|> is read as part of -/// what it is told; "start_from" (img2img: see strength) and "control" (a ControlNet's condition: -/// see control_scale) are kept for images that are drawn with, and the prompt cannot name them. -#[repr(C)] -pub struct WaifuImageInput { - /// Letters, digits and _. - pub key: *const c_char, - /// Any size, scaled by the engine. - pub image: *const WaifuImage, -} - -/// Refused at the call, as WAIFU_ERR_INVALID_ARGUMENT: two images with one key; a <|name|> in the -/// prompt with no image by that name, or naming a kept one; an image the prompt does not name and -/// whose key is not a kept one; and any other <|...|>, which the model would read as one of its -/// own markers. -#[repr(C)] -pub struct WaifuDrawRequest { - /// sizeof(WaifuDrawRequest), as the caller was compiled with. - pub struct_size: u32, - pub prompt: *const c_char, - /// Words only; "" or NULL for none; ignored by a model without guidance. - pub negative: *const c_char, - /// NULL, 0 for none. - pub images: *const WaifuImageInput, - pub image_count: usize, - /// 64..2048, multiples of the model's `alignment` in its description (16; 32 for SDXL). - pub width: u32, - pub height: u32, - pub steps: u32, - pub guidance: f32, - /// The caller's: the same seed and settings draw the same image. - pub seed: u64, - /// With "start_from": how far it walks away from it, 0..1. - pub strength: f32, - /// With "control": how hard it holds to it. - pub control_scale: f32, -} - -#[repr(C)] -pub struct WaifuImageCompletionEvent { - pub status: WaifuStatus, - /// NULL unless status.code is WAIFU_OK. - pub image: *const WaifuImage, -} - -pub type WaifuImageCompleteCallback = - Option; - /// A request struct, read only if it is at least as large as the one this library knows. unsafe fn request_in<'a, T>(request: *const T, size: u32, what: &str) -> Result<&'a T, String> { let request = request.as_ref().ok_or_else(|| format!("the {what} is NULL"))?; @@ -883,54 +540,6 @@ pub unsafe extern "C" fn waifu_engine_draw_async( // -- audio ----------------------------------------------------------------------------------- -/// The samples, not a file: what goes in -- a recording to sound like, or to convert -- and what -/// comes out. Decoding a file to this and writing this to a file are the caller's. -#[repr(C)] -pub struct WaifuAudio { - /// Mono, -1..1. - pub samples: *const f32, - pub count: usize, - pub rate: u32, -} - -#[repr(C)] -pub struct WaifuSpeakRequest { - /// sizeof(WaifuSpeakRequest), as the caller was compiled with. - pub struct_size: u32, - pub text: *const c_char, - /// A recording to sound like, or NULL. - pub like: *const WaifuAudio, - pub speed: f32, - pub temperature: f32, - /// NULL or "" for the voice's own way. - pub style: *const c_char, - pub seed: u64, -} - -#[repr(C)] -pub struct WaifuVoiceConversionRequest { - /// sizeof(WaifuVoiceConversionRequest), as the caller was compiled with. - pub struct_size: u32, - /// What is said. - pub source: *const WaifuAudio, - /// Whose voice to say it in. - pub reference: *const WaifuAudio, - pub steps: u32, - pub convert_style: bool, - pub seed: u64, -} - -/// Shared by speak and voice_conversion: both make audio. -#[repr(C)] -pub struct WaifuAudioCompletionEvent { - pub status: WaifuStatus, - /// NULL unless status.code is WAIFU_OK. - pub audio: *const WaifuAudio, -} - -pub type WaifuAudioCompleteCallback = - Option; - unsafe fn audio_in(audio: *const WaifuAudio, what: &str) -> Result { let audio = audio.as_ref().ok_or_else(|| format!("{what} is NULL"))?; if audio.samples.is_null() && audio.count > 0 { diff --git a/waifu/src/c_api/support.rs b/waifu/src/c_api/support.rs new file mode 100644 index 0000000..27ab760 --- /dev/null +++ b/waifu/src/c_api/support.rs @@ -0,0 +1,310 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! What the C API is made of besides its functions: the types every part of it takes and hands +//! back -- status, the log, progress and completion -- and the helpers that convert to and from +//! them. Written into waifu.h with the functions. + +#![allow(non_camel_case_types)] +#![allow(clippy::missing_safety_doc)] + +use std::cell::RefCell; +use std::ffi::{c_char, c_void, CStr, CString}; +use std::panic::{self, AssertUnwindSafe}; +use std::ptr; +use std::sync::{Arc, Condvar, Mutex}; + +use crate::log::{Level, Sink}; + +// -- status ---------------------------------------------------------------------------------- + +/// What a call or a job came to. +#[repr(C)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum WaifuStatusCode { + WAIFU_OK = 0, + WAIFU_CANCELLED = 1, + /// A NULL, a string that is not UTF-8, a size out of range, a request that names an image + /// it does not give. + WAIFU_ERR_INVALID_ARGUMENT = 2, + /// Neither a published name nor a manifest that exists. + WAIFU_ERR_UNKNOWN_MODEL = 3, + /// The download failed. + WAIFU_ERR_FETCH = 4, + /// The model could not be read, or cannot do what was asked. + WAIFU_ERR_MODEL = 5, + /// A job with no model loaded, or a model of the wrong kind. + WAIFU_ERR_NOT_LOADED = 6, + /// A panic in the library: a bug. + WAIFU_ERR_INTERNAL = 7, +} + +/// How a job ended, as on_complete is told it. +#[repr(C)] +pub struct WaifuStatus { + pub code: WaifuStatusCode, + /// Why, when code is not WAIFU_OK; NULL when it is. + pub message: *const c_char, +} + +thread_local! { + static LAST_ERROR: RefCell = RefCell::new(CString::default()); +} + +/// Says why this thread's last call failed, for its `_last_error` to hand back. +pub fn set_last_error(message: &str) { + let message = c_string(message); + LAST_ERROR.with(|last| *last.borrow_mut() = message); +} + +/// This thread's last message; valid until this thread's next call. +pub fn last_error() -> *const c_char { + LAST_ERROR.with(|last| last.borrow().as_ptr()) +} + +/// Frees a string [`string_out`] made. NULL is fine. +pub unsafe fn free_string(s: *mut c_char) { + if !s.is_null() { + drop(CString::from_raw(s)); + } +} + +/// Runs a synchronous entry point, turning a panic into its message in the last error and +/// `otherwise` as what the call returns. +pub fn caught(otherwise: T, body: impl FnOnce() -> T) -> T { + match panic::catch_unwind(AssertUnwindSafe(body)) { + Ok(value) => value, + Err(cause) => { + let said = cause + .downcast_ref::<&str>() + .map(|said| said.to_string()) + .or_else(|| cause.downcast_ref::().cloned()) + .unwrap_or_else(|| "it panicked".to_string()); + set_last_error(&format!("a bug in the library: {said}")); + otherwise + } + } +} + +/// A string handed in, or the reason it cannot be read. NULL is the caller's to decide. +pub unsafe fn string_in(pointer: *const c_char, what: &str) -> Result, String> { + if pointer.is_null() { + return Ok(None); + } + CStr::from_ptr(pointer) + .to_str() + .map(|text| Some(text.to_string())) + .map_err(|_| format!("{what} is not UTF-8")) +} + +/// A string handed in that may not be NULL. +pub unsafe fn required(pointer: *const c_char, what: &str) -> Result { + string_in(pointer, what)?.ok_or_else(|| format!("{what} is NULL")) +} + +/// A string for the caller to free with the library's `_string_free`. +pub fn string_out(text: &str) -> *mut c_char { + c_string(text).into_raw() +} + +/// `text` as C can read it: a NUL inside it would end it early, so it is a space instead. +pub fn c_string(text: &str) -> CString { + CString::new(text.replace('\0', " ")).unwrap_or_default() +} + +// -- logging --------------------------------------------------------------------------------- + +#[repr(C)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum WaifuLogLevel { + WAIFU_LOG_DEBUG = 0, + WAIFU_LOG_INFO = 1, + WAIFU_LOG_WARNING = 2, + WAIFU_LOG_ERROR = 3, + /// The last line: the process aborts when the callback returns. + WAIFU_LOG_FATAL = 4, +} + +#[repr(C)] +pub struct WaifuLogEvent { + pub level: WaifuLogLevel, + /// Where it was written: "interface.cc:84" for flint, a module's name for Rust. + pub source: *const c_char, + /// The line alone, with no level or time in front of it. + pub message: *const c_char, +} + +/// Called on whichever thread wrote the line, which can be the caller's own. It must be safe to +/// call from any thread, and must not call into the library. +pub type WaifuLogCallback = Option; + +/// A caller's pointer, handed back as it was given, on whichever thread a callback runs. +#[derive(Clone, Copy)] +pub struct UserData(pub *mut c_void); +unsafe impl Send for UserData {} +unsafe impl Sync for UserData {} + +/// The sink that calls `callback`, or None -- the console -- where there is no callback. +pub fn log_sink(user_data: *mut c_void, callback: WaifuLogCallback) -> Option { + let callback = callback?; + let user_data = UserData(user_data); + Some(Box::new(move |level, source, message| { + let source = c_string(source); + let message = c_string(message); + let event = WaifuLogEvent { + level: match level { + Level::Debug => WaifuLogLevel::WAIFU_LOG_DEBUG, + Level::Info => WaifuLogLevel::WAIFU_LOG_INFO, + Level::Warning => WaifuLogLevel::WAIFU_LOG_WARNING, + Level::Error => WaifuLogLevel::WAIFU_LOG_ERROR, + Level::Fatal => WaifuLogLevel::WAIFU_LOG_FATAL, + }, + source: source.as_ptr(), + message: message.as_ptr(), + }; + let user_data = user_data; + unsafe { callback(user_data.0, &event) }; + })) +} + +pub fn log_level(level: WaifuLogLevel) -> Level { + match level { + WaifuLogLevel::WAIFU_LOG_DEBUG => Level::Debug, + WaifuLogLevel::WAIFU_LOG_INFO => Level::Info, + WaifuLogLevel::WAIFU_LOG_WARNING => Level::Warning, + WaifuLogLevel::WAIFU_LOG_ERROR => Level::Error, + WaifuLogLevel::WAIFU_LOG_FATAL => Level::Fatal, + } +} + +// -- progress -------------------------------------------------------------------------------- + +#[repr(C)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum WaifuStage { + /// Reading weights onto the device: no fraction. 0 was fetching, which is + /// libwaifu_modelmanager's now; the others keep their numbers. + WAIFU_STAGE_READING = 1, + /// The prompt, or the picture img2img starts from. + WAIFU_STAGE_ENCODING = 2, + /// done/total in steps. + WAIFU_STAGE_DRAWING = 3, + WAIFU_STAGE_DECODING = 4, + /// A conversion reading its recordings. + WAIFU_STAGE_LISTENING = 5, + /// done/total in tokens; total is an estimate and can be passed. + WAIFU_STAGE_SAYING = 6, + /// The vocoder. + WAIFU_STAGE_SOUNDING = 7, +} + +#[repr(C)] +pub struct WaifuProgressEvent { + pub stage: WaifuStage, + /// 0..1 of the whole job, or -1 where nothing can say. + pub fraction: f64, + /// What the stage counts; 0, 0 where it counts nothing. + pub done: u64, + pub total: u64, + /// Reserved: 0, 0 and NULL. They said which file of a download, which is + /// libwaifu_modelmanager's now, and are kept so the event keeps its layout. + pub part: u32, + pub parts: u32, + pub file: *const c_char, + /// "step 3 of 8", for a status line. + pub words: *const c_char, + /// Since the job started. + pub seconds: f64, +} + +pub type WaifuProgressCallback = Option; + +/// How a job that makes nothing ended: a fetch, a load, an unload. +#[repr(C)] +pub struct WaifuCompletionEvent { + pub status: WaifuStatus, +} + +pub type WaifuCompleteCallback = Option; + +/// A progress report, before it is made into C. +pub struct Progress<'a> { + pub stage: WaifuStage, + pub fraction: Option, + pub done: u64, + pub total: u64, + pub part: u32, + pub parts: u32, + pub file: Option<&'a str>, + pub words: &'a str, + pub seconds: f64, +} + +/// Calls `callback`, if there is one, with `progress`. +pub fn tell_progress(callback: WaifuProgressCallback, user_data: UserData, progress: &Progress) { + let Some(callback) = callback else { + return; + }; + let file = progress.file.map(c_string); + let words = c_string(progress.words); + let event = WaifuProgressEvent { + stage: progress.stage, + fraction: progress.fraction.unwrap_or(-1.0), + done: progress.done, + total: progress.total, + part: progress.part, + parts: progress.parts, + file: file.as_ref().map_or(ptr::null(), |file| file.as_ptr()), + words: words.as_ptr(), + seconds: progress.seconds, + }; + unsafe { callback(user_data.0, &event) }; +} + +/// A status whose message lives as long as `message` does. +pub fn status(code: WaifuStatusCode, message: &Option) -> WaifuStatus { + WaifuStatus { + code, + message: message.as_ref().map_or(ptr::null(), |message| message.as_ptr()), + } +} + +// -- keeping the promises about when a callback can arrive --------------------------------------- + +/// Shut until the `_async` call that made it has returned: a callback that arrives before then +/// waits. A job's thread can finish before the call that started it is done handing back its +/// handle, and a caller whose continuation is set up from that handle would miss the end. +#[derive(Clone, Default)] +pub struct Gate(Arc<(Mutex, Condvar)>); + +impl Gate { + pub fn open(&self) { + let (open, opened) = &*self.0; + *open.lock().unwrap_or_else(|held| held.into_inner()) = true; + opened.notify_all(); + } + + pub fn wait(&self) { + let (open, opened) = &*self.0; + let mut is_open = open.lock().unwrap_or_else(|held| held.into_inner()); + while !*is_open { + is_open = opened.wait(is_open).unwrap_or_else(|held| held.into_inner()); + } + } +} diff --git a/waifu-ffi/src/tests.rs b/waifu/src/c_api/tests.rs similarity index 90% rename from waifu-ffi/src/tests.rs rename to waifu/src/c_api/tests.rs index afa0c01..f862080 100644 --- a/waifu-ffi/src/tests.rs +++ b/waifu/src/c_api/tests.rs @@ -22,9 +22,12 @@ use super::*; +use std::ffi::CStr; use std::sync::mpsc::{self, Receiver, Sender}; +use std::sync::Mutex; use std::time::Duration; + const WAIT: Duration = Duration::from_secs(20); /// What a test hears of one job, sent from the callbacks. @@ -108,18 +111,14 @@ fn the_abi_is_the_header_s() { } #[test] -fn what_is_published_and_what_a_model_is_are_json_and_freed_by_the_caller() { +fn what_a_model_is_and_the_machine_are_json_and_freed_by_the_caller() { unsafe { - let catalog = waifu_modelmanager_catalog_json(); - assert!(!catalog.is_null()); - assert!(text(catalog).unwrap().contains("\"sdxl:base\"")); - waifu_string_free(catalog); - - let described = waifu_modelmanager_describe_json(c"sdxl:base".as_ptr()); + let manifest = std::ffi::CString::new(crate::engine::tests::a_manifest("c-api-sdxl", "sdxl")).unwrap(); + let described = waifu_describe_json(manifest.as_ptr()); assert!(text(described).unwrap().contains("\"image_keys\"")); waifu_string_free(described); - assert!(waifu_modelmanager_describe_json(c"no-such-model:v9".as_ptr()).is_null()); + assert!(waifu_describe_json(c"no-such-model:v9".as_ptr()).is_null()); assert!(last_error().contains("no-such-model:v9"), "{}", last_error()); let machine = waifu_machine_json(); @@ -128,18 +127,6 @@ fn what_is_published_and_what_a_model_is_are_json_and_freed_by_the_caller() { } } -#[test] -fn bundled_manifests_are_a_directory_or_none() { - unsafe { - // Where they are is the caller's; what is read from them is hub's to test. - let directory = std::env::temp_dir(); - let directory = std::ffi::CString::new(directory.to_str().expect("a UTF-8 path")).unwrap(); - assert_eq!(waifu_modelmanager_set_bundled_manifests(directory.as_ptr()), WAIFU_OK); - assert_eq!(waifu_modelmanager_set_bundled_manifests(c"".as_ptr()), WAIFU_OK); - assert_eq!(waifu_modelmanager_set_bundled_manifests(ptr::null()), WAIFU_OK); - } -} - #[test] fn a_call_refused_at_the_door_returns_0_and_calls_nothing() { unsafe { @@ -253,7 +240,7 @@ fn freeing_the_engine_ends_every_job_it_had() { } #[test] -fn a_model_nobody_published_is_unknown_at_the_end() { +fn a_name_that_is_no_manifest_is_unknown_at_the_end() { unsafe { let engine = waifu_engine_new(WaifuDevice::WAIFU_DEVICE_CPU); let (send, heard) = a_channel(); @@ -281,7 +268,7 @@ unsafe extern "C" fn on_log(user_data: *mut c_void, event: *const WaifuLogEvent) fn a_log_callback_hears_the_library_s_lines() { let lines: Box>> = Box::default(); waifu_set_log_callback(&*lines as *const _ as *mut c_void, Some(on_log)); - waifu::log::write(waifu::log::Level::Warning, "test", "a line for the app"); + crate::log::write(crate::log::Level::Warning, "test", "a line for the app"); waifu_set_log_callback(ptr::null_mut(), None); let lines = lines.lock().unwrap(); diff --git a/waifu/src/cosyvoice3/frontend.rs b/waifu/src/cosyvoice3/frontend.rs index b023b19..24dc85c 100644 --- a/waifu/src/cosyvoice3/frontend.rs +++ b/waifu/src/cosyvoice3/frontend.rs @@ -21,7 +21,7 @@ //! `frontend_utils.py` it calls, in the order it calls them. //! //! Upstream normalizes numbers with WeTextProcessing when it is installed, and so does this -- -//! through [`crate::indextts::normalize`], which is the same grammar's behaviour written as rules, +//! through [`crate::text_normalizer`], which is the same grammar's behaviour written as rules, //! with the differences that module lists. Everything else is ported as written: a Chinese text //! loses its blanks between Chinese characters, its superscripts, its full-width brackets and its //! trailing commas, and has its `.` made `。`; either language is then cut at sentence punctuation @@ -30,7 +30,7 @@ //! A text holding `<|` and `|>` -- one that carries its own markers -- is not touched, as upstream //! does not touch it. -use crate::indextts::normalize::{normalize, Language}; +use crate::text_normalizer::{normalize, Language}; use crate::Result; const TOKEN_MAX: usize = 80; diff --git a/waifu/src/describe.rs b/waifu/src/describe.rs index 2866b32..f8808cf 100644 --- a/waifu/src/describe.rs +++ b/waifu/src/describe.rs @@ -17,43 +17,36 @@ // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -//! What a model is, said without reading its weights: what a window offers before a run. +//! What a model is, said without reading its weights: what a request for it may hold. +//! +//! What a screen should start the boxes at -- a kind's defaults, and what a manifest's +//! `suggested:` block says -- is not here: running a model needs none of it, and it is the model +//! manager's to say (`waifu_modelmanager_defaults_json`). use std::path::{Path, PathBuf}; use crate::cosyvoice3::{self, CosyVoice3}; -use crate::hub; use crate::indextts::{self, IndexTts}; -use crate::{ - Anima, ConversionDefaults, GenerationDefaults, Krea2, Manifest, QwenImage, Sdxl, - SpeechDefaults, Style, Suggestions, Tones, Voice, -}; +use crate::{Anima, Krea2, Manifest, QwenImage, Sdxl, Style, Tones, Voice}; /// The model the page is set to draw with, as the screen describes it. /// /// Chosen is not loaded. Picking one is a click and reading one is minutes, so what is here is -/// what could be known when it was picked -- guessed from the name, for a model that may not -/// even be on the disk yet -- and it is described again out of the package the first time a run -/// needs the weights. Everything on the screen that is about the model reads this: the boxes it -/// fills in, and whether image to image is offered at all. +/// what could be known when it was picked -- read from its manifest, whose weights may not even +/// be on the disk yet -- and it is described again out of the package the first time a run needs +/// the weights. #[derive(Clone)] pub struct Chosen { - /// What to ask for when the time comes to read it: the catalogue name, or the path that was - /// given. Not for showing -- a path does not fit on the screen -- but it is what a run is - /// posted with, and what the worker hands to the hub. + /// What to ask for when the time comes to read it: the path of its manifest. Not for + /// showing -- a path does not fit on the screen -- but it is what a load is asked with. pub name: String, - /// What it is called on screen: the catalogue's full name, or the file name for a path. + /// What it is called on screen: the manifest's file name. pub full_name: String, - /// Whether every package of it is already fetched. A model that is not is still a model that + /// Whether every file its manifest names is beside it. A model that is not is still a model that /// can be chosen; the download happens at the first run, where the bar can say so. pub on_disk: bool, /// Whether its weights are read and on the device. False until the first run needs them. pub in_memory: bool, - /// What the boxes start at. Guessed from the name until the package is read, and read out of - /// the package after that -- a distilled release answers both of these differently. - pub defaults: GenerationDefaults, - /// The sizes to offer: the model's own where its manifest names enough of them. - pub sizes: Vec<(i32, i32)>, /// What both sides have to be a multiple of, which is what the kind's halvings need: a size /// typed in is rounded to this, and one that is not is refused by the run. pub alignment: i32, @@ -62,26 +55,17 @@ pub struct Chosen { /// out unlike another tool's is a question about this first, and a screen that does not say /// is a screen that cannot be asked. pub sampler: &'static str, - /// What the package's own card suggests be typed in the boxes. Known once it has been read, - /// which is why the boxes take it whenever it arrives rather than only when a page opens. - pub suggested_prompt: Option, - pub suggested_avoid: Option, /// Why this one cannot be handed a picture to start from, or None where it can. A page that /// offers what the model will refuse is a refusal someone finds out about after waiting, and /// one that is simply greyed out is a screen that cannot be asked why. pub no_picture_because: Option, - /// Whether to offer guidance and a negative prompt at all. False for a distilled release, - /// which has no second pass for either to reach. Guessed from the name like the rest of this - /// and answered by the package once it has been read. - pub takes_guidance: bool, } /// The voice the page is set to speak with, as the screen describes it. /// /// What [`Chosen`] is for a picture model, and for the same reasons: it is what could be known /// before any weights were read, and it is described again out of the voice itself once one is -/// loaded. There is one of these even before anything has been asked for, because the boxes on -/// the speech tab are a voice's numbers and have to start somewhere. +/// loaded. #[derive(Clone)] pub struct Spoken { /// What to ask for when the time comes to read it. It travels with a run the way a model's @@ -95,8 +79,6 @@ pub struct Spoken { /// Whether it is read and on the device. The stand-in, `tones`, holds nothing, so for it this /// is true from the first run rather than after a fetch. pub in_memory: bool, - /// What the boxes start at, which is the voice's to say. - pub defaults: SpeechDefaults, /// The rate it writes at, which is a property of the voice and not a setting. On the screen /// because a clip that came out at the wrong pitch is a question about this first. pub rate: u32, @@ -122,8 +104,6 @@ pub struct ChosenConverter { pub full_name: String, pub on_disk: bool, pub in_memory: bool, - /// What the boxes start at. - pub defaults: ConversionDefaults, /// The rate it writes at. pub rate: u32, /// Why it cannot convert the style as well, or None where it can -- which is where the page @@ -131,20 +111,6 @@ pub struct ChosenConverter { pub no_style_because: Option, } -/// The sizes offered for a model that does not name its own, which are the ones SDXL was trained -/// at. Every one of them is a multiple of the 32 pixels the U-Net's own halvings need. -pub const SIZES: &[(i32, i32)] = &[ - (512, 512), - (640, 640), - (768, 768), - (896, 896), - (1024, 1024), - (832, 1216), - (1216, 832), - (768, 1344), - (1344, 768), -]; - /// Why Anima cannot be handed a picture to start from. True of the kind rather than of any one /// package, which is what lets it be said before a package has been looked at. pub const ANIMA_DRAWS_FROM_NO_PICTURE: &str = "Anima cannot start from a picture yet: its package \ @@ -167,49 +133,29 @@ pub const QWEN_IMAGE_DRAWS_FROM_NO_PICTURE: &str = "Qwen-Image 2.1 cannot start /// noise where the syllables are, which is for checking the page and not for listening to. pub const TONES: &str = "tones"; -/// How many sizes a model has to name before its list is used instead of the one above. -/// -/// One size is not a choice and two is barely one. A model that names fewer has not replaced the -/// list, it has emptied it, and the built-in one is the better thing to offer. -pub const ENOUGH_SIZES: usize = 3; - /// The voice `asked` names, as the screen describes it, without reading any of it. /// /// What [`look_at`] is for a picture model. [`TONES`] is [`Tones`], built into the binary, which -/// can be constructed and asked about itself for nothing. Anything else is a package -- a manifest -/// on the disk, or a published name -- and what can be said of one before minutes of reading is -/// what its kind says: IndexTTS-2.5's rate and starting values, whether it is on the disk, and -/// that it is not in memory. The read itself is [`read_voice`], at the first run that wants it, -/// and it says so if the package turns out to be something else. +/// can be constructed and asked about itself for nothing. Anything else is the path of a manifest, +/// and what can be said of one before minutes of reading is what its kind says: IndexTTS-2.5's +/// rate, whether its files are on the disk, and that it is not in memory. pub fn look_at_voice(asked: &str) -> Spoken { if asked == TONES { return describe_voice(asked, &Tones::new(), false); } - let path = manifest_of(asked); - // The manifest says what it is where there is one; a published voice not yet fetched is - // known by its family, so the page offers its own settings before the download. - let kind = path - .as_deref() - .and_then(voice_kind) - .or_else(|| published_voice_kind(asked)); - let (kind_name, defaults, rate, styles) = match kind { - Some(VoiceKind::CosyVoice3) => ( - CosyVoice3::NAME, - CosyVoice3::DEFAULTS, - cosyvoice3::RATE, - CosyVoice3::STYLES, - ), - _ => (IndexTts::NAME, IndexTts::DEFAULTS, indextts::RATE, &[][..]), + let kind = manifest_of(asked).as_deref().and_then(voice_kind); + let (kind_name, rate, styles) = match kind { + Some(VoiceKind::CosyVoice3) => (CosyVoice3::NAME, cosyvoice3::RATE, CosyVoice3::STYLES), + _ => (IndexTts::NAME, indextts::RATE, &[][..]), }; Spoken { name: asked.to_string(), - // The catalogue's name for it where it has one, and the kind's where it was named by path. - full_name: hub::full_name(asked).unwrap_or(kind_name).to_string(), + // The kind's name: what the catalogue calls a published one is the model manager's. + full_name: kind_name.to_string(), on_disk: on_disk(asked).is_some(), in_memory: false, - defaults, rate, no_likeness_because: None, styles, @@ -217,12 +163,10 @@ pub fn look_at_voice(asked: &str) -> Spoken { } } -/// Whether `asked` is a voice: the stand-in, a published voice, or a manifest on the disk whose -/// `model.type` is one of the speech models. What `-m` alone is taken to mean text2speech by. +/// Whether `asked` is a voice: the stand-in, or a manifest on the disk whose `model.type` is one +/// of the speech models. pub fn is_a_voice(asked: &str) -> bool { - asked == TONES - || hub::is_voice(asked) - || manifest_of(asked).as_deref().and_then(voice_kind).is_some() + asked == TONES || manifest_of(asked).as_deref().and_then(voice_kind).is_some() } /// The speech models a package can be. @@ -232,15 +176,6 @@ pub enum VoiceKind { CosyVoice3, } -/// Which speech model a published voice is, by the family its versioned name starts with. -pub fn published_voice_kind(asked: &str) -> Option { - match hub::published_name(asked)?.split(':').next()? { - "indextts" => Some(VoiceKind::IndexTts), - "cosyvoice" => Some(VoiceKind::CosyVoice3), - _ => None, - } -} - /// Which speech model the manifest at `path` describes, by its `model.type`; `None` where it /// cannot be read or is not a speech model this reads. pub fn voice_kind(path: &std::path::Path) -> Option { @@ -256,12 +191,10 @@ pub fn voice_kind(path: &std::path::Path) -> Option { pub fn describe_voice(name: &str, voice: &dyn Voice, in_memory: bool) -> Spoken { Spoken { name: name.to_string(), - // The catalogue's name first, so the button does not change its wording when it is read. - full_name: hub::full_name(name).unwrap_or(voice.name()).to_string(), + full_name: voice.name().to_string(), // Built in, or just read off the disk: either way there is nothing left to fetch. on_disk: true, in_memory, - defaults: voice.defaults(), rate: voice.rate(), no_likeness_because: voice.no_likeness_because().map(str::to_string), styles: voice.styles(), @@ -275,20 +208,18 @@ pub fn look_at_converter(asked: &str) -> ChosenConverter { // CosyVoice3 is the one converter there is, so there is no kind to tell apart yet. ChosenConverter { name: asked.to_string(), - full_name: hub::full_name(asked).unwrap_or(CosyVoice3::NAME).to_string(), + full_name: CosyVoice3::NAME.to_string(), on_disk: on_disk(asked).is_some(), in_memory: false, - defaults: CosyVoice3::CONVERSION_DEFAULTS, rate: cosyvoice3::RATE, no_style_because: Some(CosyVoice3::NO_STYLE_CONVERSION.to_string()), } } -/// Whether `asked` turns one recording into another: a published converter, or a manifest on the -/// disk whose `model.type` is one. What `-m` alone is taken to mean speech2speech by, after -/// [`is_a_voice`] -- CosyVoice3 is both, and alone it reads. +/// Whether `asked` turns one recording into another: a manifest on the disk whose `model.type` is +/// one. CosyVoice3 is both this and a voice. pub fn is_a_converter(asked: &str) -> bool { - hub::is_conversion(asked) || manifest_of(asked).as_deref().and_then(converter_kind).is_some() + manifest_of(asked).as_deref().and_then(converter_kind).is_some() } /// The models a package can convert voices with. @@ -306,41 +237,16 @@ pub fn converter_kind(path: &Path) -> Option { } } -/// What a kind of model implies for the screen before any of its weights are read: what to ask -/// it for, what walks the noise back, why it cannot start from a picture, and whether guidance is -/// a thing to offer at all. -/// -/// One table for the three kinds, read twice -- once from a name that has not been fetched yet, -/// and once from the `type` of a manifest that is already on the disk. -/// -/// The last of the four is false for Krea 2 for the same reason its defaults are eight steps at a -/// guidance of one: the only release of it this exports is the distilled one, and the distilled -/// one has no second pass. A package's own `suggested:` overrides this, so the day there is an -/// undistilled one it says `takes_guidance: "true"` and gets the dial back. -pub fn about(kind: &str) -> (GenerationDefaults, &'static str, Option<&'static str>, bool) { +/// What a kind of model implies before any of its weights are read: what walks the noise back, +/// and why it cannot start from a picture. One table for the kinds, read from the `type` a +/// manifest states. +pub fn about(kind: &str) -> (&'static str, Option<&'static str>) { match kind { - Anima::MODEL_TYPE => ( - Anima::DEFAULTS, - "Flow match Euler", - Some(ANIMA_DRAWS_FROM_NO_PICTURE), - true, - ), - Krea2::MODEL_TYPE => ( - Krea2::DEFAULTS, - "Flow match Euler", - Some(KREA2_DRAWS_FROM_NO_PICTURE), - false, - ), - // Sampled without guidance by default, and able to take it: the reference's - // `true_cfg_scale` is a second pass on a negative prompt, which is this runtime's. - QwenImage::MODEL_TYPE => ( - QwenImage::DEFAULTS, - "Flow match Euler", - Some(QWEN_IMAGE_DRAWS_FROM_NO_PICTURE), - true, - ), + Anima::MODEL_TYPE => ("Flow match Euler", Some(ANIMA_DRAWS_FROM_NO_PICTURE)), + Krea2::MODEL_TYPE => ("Flow match Euler", Some(KREA2_DRAWS_FROM_NO_PICTURE)), + QwenImage::MODEL_TYPE => ("Flow match Euler", Some(QWEN_IMAGE_DRAWS_FROM_NO_PICTURE)), // Everything else is SDXL, which is what `Model::from_manifest` decides too. - _ => (GenerationDefaults::default(), "Euler", None, true), + _ => ("Euler", None), } } @@ -354,55 +260,32 @@ pub fn alignment(kind: &str) -> i32 { } } -/// What can be said about a model that has not been read, which is what a choice is made from. +/// What can be said about a model that has not been read, which is what a choice is made from: +/// what the manifest at `asked` says, and none of its weights. Not the last word -- the moment a +/// run reads the package, the engine describes it again out of what is actually in there. /// -/// Guessed from the name where there is nothing else to go on, and read out of the manifest where -/// the package is already here. Neither is the last word -- the moment a run reads the package, -/// [`load`] describes it again out of what is actually in there, wrong guesses included. +/// The manifest need not have its packages beside it: one of the bundled ones, or one fetched on +/// its own, says what the model is before gigabytes of it are downloaded. A manifest that cannot +/// be read describes an SDXL model, because the one thing this decides -- whether image to image +/// is offered -- is a thing SDXL packages do. pub fn look_at(asked: &str) -> Chosen { - // The kinds are told apart by the catalogue's own naming. A package somebody exported - // themselves and named by path is taken as SDXL until the manifest below says otherwise, - // because the one thing this decides -- whether image to image is offered -- is a thing SDXL - // packages do. - let guessed = match asked { - name if name.starts_with("anima") => Anima::MODEL_TYPE, - name if name.starts_with("krea") => Krea2::MODEL_TYPE, - name if name.starts_with("qwen") => QwenImage::MODEL_TYPE, - _ => "", - }; - let (defaults, sampler, no_picture, takes_guidance) = about(guessed); + let (sampler, no_picture) = about(""); let mut chosen = Chosen { name: asked.to_string(), - // The catalogue's name for it, or the file name where a path was given: the whole path - // does not fit on screen and its tail is the part that identifies it. - full_name: hub::full_name(asked) - .map(str::to_string) - .unwrap_or_else(|| match asked.rsplit_once('/') { - Some((_, file)) if !file.is_empty() => file.to_string(), - _ => asked.to_string(), - }), - on_disk: hub::is_cached(asked), + // The file name: the whole path does not fit on screen and its tail is the part that + // identifies it. What the catalogue calls a published one is the model manager's. + full_name: shown_name(asked), + on_disk: on_disk(asked).is_some(), in_memory: false, - defaults, - sizes: SIZES.to_vec(), - alignment: alignment(guessed), + alignment: alignment(""), sampler, - // Not known until the package is read: what a card suggests be typed is written in it. - suggested_prompt: None, - suggested_avoid: None, no_picture_because: no_picture.map(str::to_string), - takes_guidance, }; - // And where its manifest is here -- with its packages or, fetched on its own, without them -- - // what it says about itself. The manifest is a few - // kilobytes of text beside gigabytes of weights, and reading it is what keeps the numbers in - // the boxes from changing under somebody later: a model described twice, once from its name - // and once from its package, is a model whose settings move while they are being used. - // - // The weights are still not touched. Whether this one can start from a picture is the guess - // above until a run opens the package. + // What the manifest says about itself. It is a few kilobytes of text beside gigabytes of + // weights, and the weights are not touched: whether this one can start from a picture is what + // its kind says until a run opens the package. let Some(path) = manifest_of(asked) else { return chosen; }; @@ -410,93 +293,56 @@ pub fn look_at(asked: &str) -> Chosen { return chosen; }; - // What kind it is, which the manifest states and a name only hints at. It matters most for a - // package named by its path, which is every package somebody exported themselves: `krea2` is - // not in the catalogue at all, so this is the only place its kind is ever read. + // What kind it is, which the manifest states. if let Ok(kind) = manifest .section(Sdxl::MODEL_SECTION) .and_then(|section| section.get_str("type").map(str::to_string)) { - let (defaults, sampler, no_picture, takes_guidance) = about(&kind); - chosen.defaults = defaults; + let (sampler, no_picture) = about(&kind); chosen.sampler = sampler; chosen.no_picture_because = no_picture.map(str::to_string); - chosen.takes_guidance = takes_guidance; chosen.alignment = alignment(&kind); } - let suggested = manifest.suggested(); - chosen.defaults = suggested.over(chosen.defaults); - chosen.suggested_prompt = suggested.prompt.clone(); - chosen.suggested_avoid = suggested.avoid.clone(); - chosen.takes_guidance = guided_by(suggested, chosen.takes_guidance); - if suggested.sizes.len() >= ENOUGH_SIZES { - chosen.sizes = suggested.sizes.clone(); - } - chosen } -/// Whether to offer guidance and a negative prompt, for a package suggesting `suggested`, of a kind -/// that offers them where `kind_does`. -/// -/// What the package says, where it says it. Where it does not, a card that asks for a guidance of -/// one has said it all the same: at one there is no unconditional pass, so a dial above it is a -/// setting the card advises against and a negative prompt is text nothing reads. Anima Turbo is -/// the one published like that -- "at CFG 1 and 8-12 steps", and no `takes_guidance` key. Silence -/// with any other number leaves the kind's answer alone. -fn guided_by(suggested: &Suggestions, kind_does: bool) -> bool { - suggested - .takes_guidance - .unwrap_or(kind_does && suggested.guidance != Some(1.0)) -} - -/// The manifest of `asked` wherever it is here, whether or not its packages are: what to read for -/// what it is and what it suggests. A published model's comes down first, and on its own when a -/// screen asks for every model's -- see [`hub::fetch_manifests`]. +/// The manifest at `asked`, where there is one. pub fn manifest_of(asked: &str) -> Option { - hub::local_manifest(asked).or_else(|| { - let path = PathBuf::from(asked); - path.is_file().then_some(path) - }) + let path = PathBuf::from(asked); + path.is_file().then_some(path) } -/// The manifest of a model that can be read without fetching anything: one already in the cache, -/// or one that was named by the path of its manifest in the first place. +/// The manifest at `asked`, where every file it names is beside it: a model that can be read +/// without fetching anything. pub fn on_disk(asked: &str) -> Option { - hub::cached_manifest(asked).or_else(|| { - let path = PathBuf::from(asked); - path.is_file().then_some(path) - }) + bytes_on_disk(asked).map(|_| PathBuf::from(asked)) } -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn a_card_that_asks_for_a_guidance_of_one_takes_none() { - let turbo = Suggestions { - guidance: Some(1.0), - ..Suggestions::default() - }; - assert!(!guided_by(&turbo, true), "Anima Turbo's card says CFG 1 and nothing else"); - - // What the package says outranks what it suggests: an undistilled model that likes one. - let says_so = Suggestions { - guidance: Some(1.0), - takes_guidance: Some(true), - ..Suggestions::default() - }; - assert!(guided_by(&says_so, true)); - - // Silence with another number, or with none, is the kind's answer. - let anima = Suggestions { - guidance: Some(4.0), - ..Suggestions::default() - }; - assert!(guided_by(&anima, true)); - assert!(guided_by(&Suggestions::default(), true)); - assert!(!guided_by(&Suggestions::default(), false)); +/// How many bytes the files the manifest at `asked` names are, where all of them are here: what +/// choosing a device it fits on weighs. +pub fn bytes_on_disk(asked: &str) -> Option { + let manifest = Manifest::open(manifest_of(asked)?).ok()?; + let mut bytes = 0; + for name in manifest.files() { + bytes += std::fs::metadata(manifest.file(&name).ok()?).ok()?.len(); } + Some(bytes) +} + +/// The name a screen shows for `asked`: the file name of a path, since the whole of one does not +/// fit, and the name itself otherwise. +pub fn shown_name(asked: &str) -> String { + match Path::new(asked).file_name().and_then(|name| name.to_str()) { + Some(file) if !file.is_empty() => file.to_string(), + _ => asked.to_string(), + } +} + +/// Why `asked` is not a model this can read. +pub fn no_manifest_at(asked: &str) -> String { + format!( + "there is no manifest at \"{asked}\": a model is named by the path of its manifest, which \ + waifu_modelmanager_manifest_path hands back for a published one" + ) } diff --git a/waifu/src/engine.rs b/waifu/src/engine.rs index 1cda947..ecf5742 100644 --- a/waifu/src/engine.rs +++ b/waifu/src/engine.rs @@ -24,7 +24,7 @@ //! that thread. The calls return at once; what comes of a job arrives through its callbacks -- //! progress any number of times, then exactly one completion -- on the engine's thread. //! -//! This is what the C API in waifu-ffi is a thin skin over; see docs/ffi.md for the contract it +//! This is what the C API in [`crate::c_api`] is a thin skin over; see docs/ffi.md for the contract it //! keeps. One model at a time: two on one card double the memory, and on Metal nothing can say //! beforehand whether they would fit. @@ -40,11 +40,10 @@ use std::time::Instant; use crate::cosyvoice3::CosyVoice3; use crate::describe::{self, converter_kind, voice_kind, Chosen, ConverterKind, VoiceKind, TONES}; -use crate::hub; use crate::image_model::ImageModel; use crate::indextts::IndexTts; -use crate::progress::{Convert, Doing, Fetch, Run, Say, VOCODER_RATIO}; -use crate::runtime::{DeviceOption, Runtime}; +use crate::progress::{Convert, Doing, Run, Say, VOCODER_RATIO}; +use crate::runtime::{DeviceOption, Room, Runtime}; use crate::wav::Sound; use crate::{ from_rgb8, to_rgb8, ConversionOptions, ConversionProgress, Converter, Device, @@ -108,7 +107,6 @@ pub struct VoiceConversionRequest { /// Which part of a job a progress report is about. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum Stage { - Fetching, Reading, Encoding, Drawing, @@ -124,13 +122,9 @@ pub struct Progress { pub stage: Stage, /// Of the whole job, 0 to 1, or None where nothing can say -- reading weights. pub fraction: Option, - /// What the stage counts: bytes, steps or tokens. 0 and 0 where it counts nothing. + /// What the stage counts: steps or tokens. 0 and 0 where it counts nothing. pub done: u64, pub total: u64, - /// Fetching only: which file of how many. - pub part: u32, - pub parts: u32, - pub file: Option, /// "step 3 of 8", for a status line. pub words: String, /// Since the job started. @@ -140,8 +134,8 @@ pub struct Progress { /// Why a job failed, in the C API's terms. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum Failure { + /// Neither the stand-in nor a manifest that exists. UnknownModel, - Fetch, Model, NotLoaded, Internal, @@ -278,7 +272,9 @@ impl Engine { job } - /// Fetches `model` if it has to, and reads it onto the device, dropping the one before. + /// Reads the model whose manifest is at `model` onto the device, dropping the one before. + /// `model` is a path -- fetching a published model and saying where it went is the model + /// manager's -- or the stand-in voice, "tones". pub fn load(&self, model: &str, on_progress: OnProgress, on_complete: OnComplete<()>) -> JobId { self.submit(Work::Load { model: model.to_string(), @@ -528,41 +524,33 @@ struct Reporter<'a> { impl Reporter<'_> { fn tell(&mut self, doing: &Doing) { - let (stage, done, total, part, parts, file) = match doing { + let (stage, done, total) = match doing { Doing::Nothing => return, - Doing::Fetching(fetch) => ( - Stage::Fetching, - fetch.done, - fetch.total.unwrap_or(0), - fetch.part as u32, - fetch.parts as u32, - (!fetch.file.is_empty()).then(|| fetch.file.clone()), - ), - Doing::Reading { .. } => (Stage::Reading, 0, 0, 0, 0, None), + Doing::Reading { .. } => (Stage::Reading, 0, 0), Doing::Drawing(run) => match run.progress { - GenerationProgress::Encoding => (Stage::Encoding, 0, 0, 0, 0, None), + GenerationProgress::Encoding => (Stage::Encoding, 0, 0), GenerationProgress::Step { done, total } => { - (Stage::Drawing, done as u64, total as u64, 0, 0, None) + (Stage::Drawing, done as u64, total as u64) } - GenerationProgress::Decoding => (Stage::Decoding, 0, 0, 0, 0, None), + GenerationProgress::Decoding => (Stage::Decoding, 0, 0), }, Doing::Speaking(say) => match say.progress { - SpeechProgress::Reading => (Stage::Encoding, 0, 0, 0, 0, None), + SpeechProgress::Reading => (Stage::Encoding, 0, 0), SpeechProgress::Saying { done, expected } => { - (Stage::Saying, done.max(0) as u64, expected.max(0) as u64, 0, 0, None) + (Stage::Saying, done.max(0) as u64, expected.max(0) as u64) } - SpeechProgress::Sounding => (Stage::Sounding, 0, 0, 0, 0, None), + SpeechProgress::Sounding => (Stage::Sounding, 0, 0), }, Doing::Converting(convert) => match convert.progress { - ConversionProgress::Listening => (Stage::Listening, 0, 0, 0, 0, None), + ConversionProgress::Listening => (Stage::Listening, 0, 0), ConversionProgress::Saying { done, expected } => { - (Stage::Saying, done.max(0) as u64, expected.max(0) as u64, 0, 0, None) + (Stage::Saying, done.max(0) as u64, expected.max(0) as u64) } ConversionProgress::Drawing { done, total } => { - (Stage::Drawing, done.max(0) as u64, total.max(0) as u64, 0, 0, None) + (Stage::Drawing, done.max(0) as u64, total.max(0) as u64) } ConversionProgress::Sounding { done, total } => { - (Stage::Sounding, done.max(0) as u64, total.max(0) as u64, 0, 0, None) + (Stage::Sounding, done.max(0) as u64, total.max(0) as u64) } }, }; @@ -572,142 +560,12 @@ impl Reporter<'_> { fraction: doing.fraction(), done, total, - part, - parts, - file, words: doing.words(), seconds: self.started.elapsed().as_secs_f64(), }); } } -/// The name a screen shows for `asked`: the file name of a path, since the whole of one does not -/// fit, and the name itself otherwise. -fn shown_name(asked: &str) -> String { - match asked.rsplit_once('/') { - Some((_, file)) if !file.is_empty() => file.to_string(), - _ => asked.to_string(), - } -} - -/// Fetches the published model `asked` names, if it is not here already, and hands back its -/// manifest; a manifest's own path is handed back as it is. On the caller's thread: what -/// `waifu_modelmanager_fetch_async` runs on a thread of its own, and what a load starts with. -pub fn fetch(asked: &str, on_progress: &mut OnProgress, stop: &dyn Fn() -> bool) -> Ended { - let mut reporter = Reporter { - on_progress, - started: Instant::now(), - }; - fetch_reporting(asked, &mut reporter, stop) -} - -/// Fetches the manifest of every published model not yet here, and none of their packages: see -/// [`hub::fetch_manifests`]. Each is reported as a part of the whole, under the model name -/// "model metadata". -pub fn fetch_manifests(on_progress: &mut OnProgress, stop: &dyn Fn() -> bool) -> Ended<()> { - let mut reporter = Reporter { - on_progress, - started: Instant::now(), - }; - let mut fetch = Fetch { - model: "model metadata".to_string(), - hub: None, - file: String::new(), - done: 0, - total: None, - part: 0, - parts: 0, - }; - let fetched = hub::fetch_manifests( - &mut |progress| { - match progress { - hub::Progress::From { hub } => fetch.hub = Some(hub), - hub::Progress::Fetching { - file, - done, - total, - part, - parts, - } => { - fetch.file = file.to_string(); - (fetch.done, fetch.total, fetch.part, fetch.parts) = (done, total, part, parts); - } - hub::Progress::Fetched { - file, - bytes, - part, - parts, - } => { - fetch.file = file.to_string(); - (fetch.done, fetch.total, fetch.part, fetch.parts) = (bytes, Some(bytes), part, parts); - } - } - reporter.tell(&Doing::Fetching(fetch.clone())); - }, - stop, - ); - match fetched { - Ok(()) => Ended::Done(()), - Err(error) if hub::stopped(&error) => Ended::Cancelled, - Err(error) => Ended::failed(Failure::Fetch, error.to_string()), - } -} - -fn fetch_reporting(asked: &str, reporter: &mut Reporter, stop: &dyn Fn() -> bool) -> Ended { - if hub::full_name(asked).is_none() && !Path::new(asked).is_file() { - return Ended::failed( - Failure::UnknownModel, - format!("there is no model called \"{asked}\": neither a published name nor a manifest that exists"), - ); - } - - let mut fetch = Fetch { - model: shown_name(asked), - hub: None, - file: String::new(), - done: 0, - total: None, - part: 0, - parts: 0, - }; - reporter.tell(&Doing::Fetching(fetch.clone())); - - let fetched = hub::resolve_reporting( - asked, - &mut |progress| { - match progress { - hub::Progress::From { hub } => fetch.hub = Some(hub), - hub::Progress::Fetching { - file, - done, - total, - part, - parts, - } => { - fetch.file = file.to_string(); - (fetch.done, fetch.total, fetch.part, fetch.parts) = (done, total, part, parts); - } - hub::Progress::Fetched { - file, - bytes, - part, - parts, - } => { - fetch.file = file.to_string(); - (fetch.done, fetch.total, fetch.part, fetch.parts) = (bytes, Some(bytes), part, parts); - } - } - reporter.tell(&Doing::Fetching(fetch.clone())); - }, - stop, - ); - match fetched { - Ok(path) => Ended::Done(path), - Err(error) if hub::stopped(&error) => Ended::Cancelled, - Err(error) => Ended::failed(Failure::Fetch, error.to_string()), - } -} - impl Worker { fn run(&mut self, shared: &Shared, work: Work) { let stop = || shared.stop.load(Ordering::Relaxed); @@ -717,7 +575,7 @@ impl Worker { mut on_progress, on_complete, } => { - let ended = guarded(|| self.load(&model, &mut on_progress, &stop)); + let ended = guarded(|| self.load(&model, &mut on_progress)); on_complete(ended); } Work::Unload { on_complete } => { @@ -755,7 +613,7 @@ impl Worker { } } - fn load(&mut self, asked: &str, on_progress: &mut OnProgress, stop: &dyn Fn() -> bool) -> Ended<()> { + fn load(&mut self, asked: &str, on_progress: &mut OnProgress) -> Ended<()> { let mut reporter = Reporter { on_progress, started: Instant::now(), @@ -769,12 +627,11 @@ impl Worker { self.name = asked.to_string(); return Ended::Done(()); } - let path = match fetch_reporting(asked, &mut reporter, stop) { - Ended::Done(path) => path, - Ended::Cancelled => return Ended::Cancelled, - Ended::Failed(failure, message) => return Ended::Failed(failure, message), - }; - let shown = shown_name(asked); + let path = std::path::PathBuf::from(asked); + if !path.is_file() { + return Ended::failed(Failure::UnknownModel, describe::no_manifest_at(asked)); + } + let shown = describe::shown_name(asked); reporter.tell(&Doing::Reading { model: shown }); let (device, residency) = (self.runtime.device(), self.runtime.residency()); @@ -845,17 +702,14 @@ impl Worker { } } - // Both or neither, and neither for a model with no second pass to steer. - let (guidance, negative) = match chosen.takes_guidance { - true => (request.guidance, request.negative.clone()), - false => (chosen.defaults.guidance_scale, String::new()), - }; + // As asked: what to start a model's boxes at, guidance included, is the caller's to have + // decided, and the request is what it decided. let options = GenerationOptions { width: request.width as i32, height: request.height as i32, num_steps: request.steps as i32, - guidance_scale: guidance, - negative_prompt: negative, + guidance_scale: request.guidance, + negative_prompt: request.negative.clone(), seed: Some(request.seed), strength: request.strength.clamp(0.0, 1.0), }; @@ -1026,41 +880,20 @@ fn scaled(image: &Image, width: u32, height: u32) -> Vec { // -- what is there, as JSON ------------------------------------------------------------------ -/// A float as a box can show it: four places, which is finer than any setting steps. -fn showable(value: f32) -> f64 { - (f64::from(value) * 10_000.0).round() / 10_000.0 -} - -/// The published models, aliases only: `waifu_modelmanager_catalog_json`. -pub fn catalog_json() -> serde_json::Value { - let entry = |listed: hub::Listed, kind: &str| { - serde_json::json!({ - "name": listed.name, - "full_name": listed.full_name, - "kind": kind, - "cached": listed.cached, - // Whether its manifest is here, with or without its packages: what it suggests can - // be read from it. See waifu_modelmanager_fetch_manifests_async. - "manifest_here": hub::local_manifest(listed.name).is_some(), - "bytes_on_disk": listed.bytes, - "explicit": listed.explicit, - }) - }; - let models: Vec = hub::listed() - .into_iter() - .map(|listed| entry(listed, "image")) - .chain(hub::listed_voices().into_iter().map(|listed| entry(listed, "speech"))) - .chain(hub::listed_conversions().into_iter().map(|listed| entry(listed, "conversion"))) - .collect(); - serde_json::json!({ "models": models }) -} - -/// What a model is, without reading it: `waifu_modelmanager_describe_json`. None for a name that -/// is neither published nor a manifest on the disk. +/// What a model is, without reading it: `waifu_describe_json` -- what a request for it may hold, +/// and not what to start a screen's boxes at, which is the model manager's. None for anything but +/// the stand-in and the path of a manifest. `bytes` is how much its files are, where they are all +/// here: null where the manifest came without them. pub fn describe_json(asked: &str) -> Option { - if asked != TONES && hub::full_name(asked).is_none() && !Path::new(asked).is_file() { + if asked != TONES && !Path::new(asked).is_file() { return None; } + let mut described = described(asked); + described["bytes"] = describe::bytes_on_disk(asked).into(); + Some(described) +} + +fn described(asked: &str) -> serde_json::Value { let mut kinds = Vec::new(); if describe::is_a_voice(asked) { @@ -1077,7 +910,6 @@ pub fn describe_json(asked: &str) -> Option { "name": converter.name, "full_name": converter.full_name, "on_disk": converter.on_disk, - "steps": converter.defaults.steps, "rate": converter.rate, "converts_style": converter.no_style_because.is_none(), "no_style_because": converter.no_style_because, @@ -1092,8 +924,6 @@ pub fn describe_json(asked: &str) -> Option { "name": voice.name, "full_name": voice.full_name, "on_disk": voice.on_disk, - "speed": showable(voice.defaults.speed), - "temperature": showable(voice.defaults.temperature), "rate": voice.rate, "takes_a_recording": voice.no_likeness_because.is_none(), "no_likeness_because": voice.no_likeness_because, @@ -1107,21 +937,20 @@ pub fn describe_json(asked: &str) -> Option { if kinds.contains(&"conversion") { described["conversion"] = conversion(); } - return Some(described); + return described; } if kinds.first() == Some(&"conversion") { let converter = describe::look_at_converter(asked); - return Some(serde_json::json!({ + return serde_json::json!({ "kind": "conversion", "kinds": kinds, "name": converter.name, "full_name": converter.full_name, "on_disk": converter.on_disk, - "steps": converter.defaults.steps, "rate": converter.rate, "converts_style": converter.no_style_because.is_none(), "no_style_because": converter.no_style_because, - })); + }); } let chosen = describe::look_at(asked); @@ -1141,26 +970,18 @@ pub fn describe_json(asked: &str) -> Option { "prompt_images".to_string(), format!("{} reads words only", chosen.full_name).into(), ); - Some(serde_json::json!({ + serde_json::json!({ "kind": "image", "kinds": ["image"], "name": chosen.name, "full_name": chosen.full_name, "on_disk": chosen.on_disk, - "width": chosen.defaults.width, - "height": chosen.defaults.height, - "steps": chosen.defaults.num_steps, - "guidance": showable(chosen.defaults.guidance_scale), - "takes_guidance": chosen.takes_guidance, "sampler": chosen.sampler, - "sizes": chosen.sizes.iter().map(|(w, h)| serde_json::json!([w, h])).collect::>(), "alignment": chosen.alignment, - "prompt": chosen.suggested_prompt, - "avoid": chosen.suggested_avoid, "prompt_images": 0, "image_keys": image_keys, "why_not": why_not, - })) + }) } /// The machine: `waifu_machine_json`. Measured for the device `auto` would choose. @@ -1168,6 +989,30 @@ pub fn machine_json() -> serde_json::Value { crate::machine::describe(DeviceOption::Auto.resolve()) } +/// Where a run can go: `waifu_devices_json`. Every device the C API names, whether this build +/// and this machine can take a run there, and -- for CUDA and Vulkan where they can -- how much of +/// the card is free now, in bytes. Which of them suits a model is the caller's to decide. +pub fn devices_json() -> serde_json::Value { + let available = Runtime::available(); + let room = Room::measure(&available); + let devices: Vec = Runtime::ALL + .iter() + .map(|runtime| { + let free = match runtime.device() { + Device::Cuda => room.cuda, + Device::Vulkan => room.vulkan, + _ => None, + }; + serde_json::json!({ + "name": runtime.name(), + "available": available.contains(runtime), + "free": free, + }) + }) + .collect(); + serde_json::json!({ "devices": devices }) +} + /// Which device a name means, as the C API numbers them. pub fn device_option(device: Device) -> DeviceOption { match device { @@ -1180,7 +1025,7 @@ pub fn device_option(device: Device) -> DeviceOption { } #[cfg(test)] -mod tests { +pub(crate) mod tests { use super::*; use std::sync::mpsc::RecvTimeoutError; @@ -1294,10 +1139,11 @@ mod tests { } #[test] - fn a_model_nobody_published_is_unknown() { + fn a_name_that_is_no_manifest_is_unknown() { + // A published name included: it is the model manager's to turn into a path. let engine = an_engine(); let (progress, complete, heard) = listen::<()>(); - engine.load("no-such-model:v9", progress, complete); + engine.load("sdxl:base", progress, complete); assert!(matches!(heard.recv_timeout(WAIT).unwrap(), Ended::Failed(Failure::UnknownModel, _))); } @@ -1375,20 +1221,29 @@ mod tests { assert!(reports.iter().all(|report| !report.words.is_empty())); } - #[test] - fn what_is_published_and_what_a_model_is_come_back_as_json() { - let catalog = catalog_json(); - let models = catalog["models"].as_array().unwrap(); - assert!(models.iter().any(|model| model["name"] == "sdxl:base" && model["kind"] == "image")); - assert!(models.iter().any(|model| model["kind"] == "speech")); + /// A manifest of `kind` in a directory of its own, naming one weights file that is not there: + /// what a manifest fetched on its own, or bundled with an app, looks like. + pub(crate) fn a_manifest(stem: &str, kind: &str) -> String { + let directory = std::env::temp_dir().join(format!("waifu-manifest-{}-{stem}", std::process::id())); + std::fs::create_dir_all(&directory).unwrap(); + let path = directory.join(format!("{stem}.yaml")); + let text = format!("weights:\n - {stem}.safetensors\nconfig:\n model:\n type: {kind}\n"); + std::fs::write(&path, text).unwrap(); + path.to_str().unwrap().to_string() + } - let sdxl = describe_json("sdxl:base").unwrap(); + #[test] + fn what_a_model_is_comes_back_as_json() { + let sdxl = describe_json(&a_manifest("sdxl-base", "sdxl")).unwrap(); assert_eq!(sdxl["kind"], "image"); - assert_eq!(sdxl["takes_guidance"], true); assert_eq!(sdxl["prompt_images"], 0); + // What to start the boxes at is the model manager's, not what a request may hold. + for default in ["steps", "guidance", "takes_guidance", "sizes", "prompt"] { + assert!(sdxl.get(default).is_none(), "{default} in {sdxl}"); + } assert_eq!(sdxl["image_keys"], serde_json::json!(["start_from"])); - let anima = describe_json("anima:turbo").unwrap(); + let anima = describe_json(&a_manifest("anima-turbo", crate::Anima::MODEL_TYPE)).unwrap(); assert_eq!(anima["image_keys"], serde_json::json!([])); assert!(anima["why_not"]["start_from"].as_str().unwrap().contains("Anima")); assert_eq!(sdxl["alignment"], 32); @@ -1397,10 +1252,16 @@ mod tests { let tones = describe_json(TONES).unwrap(); assert_eq!(tones["kind"], "speech"); - let cosyvoice = describe_json("cosyvoice").unwrap(); + let cosyvoice = describe_json(&a_manifest("cosyvoice3", CosyVoice3::MODEL_TYPE)).unwrap(); assert_eq!(cosyvoice["kinds"], serde_json::json!(["speech", "conversion"])); - assert!(cosyvoice["conversion"]["steps"].as_i64().unwrap() > 0, "{cosyvoice}"); + assert!(cosyvoice["conversion"]["rate"].as_u64().unwrap() > 0, "{cosyvoice}"); + + // Its weights are not beside it, which is a manifest that says what it is all the same. + assert_eq!(sdxl["on_disk"], false); + assert_eq!(sdxl["bytes"], serde_json::Value::Null); + // A published name is the model manager's to turn into a path. + assert!(describe_json("sdxl:base").is_none()); assert!(describe_json("no-such-model:v9").is_none()); } } diff --git a/waifu/src/indextts/mod.rs b/waifu/src/indextts/mod.rs index ab1296e..b2a6edc 100644 --- a/waifu/src/indextts/mod.rs +++ b/waifu/src/indextts/mod.rs @@ -94,8 +94,6 @@ pub mod emotion; pub mod features; /// The GPT: text and a voice in, the semantic tokens [`crate::indextts::s2mel`] reads out. pub mod gpt; -/// The text frontend: writing out what a number is read as, so the model never sees a digit. -pub mod normalize; /// S2Mel's denoiser: semantic tokens and a voice in, the mel spectrogram a vocoder reads out. pub mod s2mel; /// The semantic codec: w2v-bert's features in, the discrete tokens the GPT reads out. @@ -119,7 +117,7 @@ use crate::{ use self::bigvgan::{BigVgan, BigVganConfig}; use self::gpt::{Gpt, Sampling}; -use self::normalize::Language; +use crate::text_normalizer::{self as normalize, Language}; /// Where each of the six models sits in the package. See `tools/indextts_package.py`. const GPT: &str = "indextts2.gpt"; diff --git a/waifu/src/lib.rs b/waifu/src/lib.rs index 6158220..ee0e4eb 100644 --- a/waifu/src/lib.rs +++ b/waifu/src/lib.rs @@ -17,12 +17,17 @@ // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -//! Diffusion image generation on top of the flint tensor library. +//! Diffusion image generation on top of the flint tensor library: libwaifu. //! //! It reads a model package, builds the model it describes, and runs it; the tensor operations //! themselves are the ones [`flint`] binds, and that module is the safe wrapper over the native //! `libflint.a` this crate links. //! +//! What anything outside this crate calls is the C API in [`c_api`], from the dynamic library -- +//! the command line too. The Rust items are public for this crate's own tests, examples and +//! benchmarks. A model is handed over as the path of its manifest: which models are published, +//! and fetching them, is libwaifu_modelmanager's. +//! //! ```no_run //! use waifu::Manifest; //! @@ -43,23 +48,13 @@ pub mod anima; /// activation, all composed from the operators [`flint`] already has rather than written as /// kernels of their own. pub mod audio; -#[cfg(feature = "cli")] -pub mod cli; -#[cfg(feature = "hub")] -pub mod config; -#[cfg(feature = "engine")] +/// The C API, which is all anything outside this crate calls: see docs/ffi.md. +pub mod c_api; pub mod describe; -#[cfg(feature = "engine")] pub mod engine; -#[cfg(feature = "hub")] -pub mod hub; -#[cfg(feature = "engine")] pub mod image_model; -#[cfg(feature = "engine")] pub mod machine; -#[cfg(feature = "engine")] pub mod progress; -#[cfg(feature = "engine")] pub mod runtime; /// Fun-CosyVoice3-0.5B: a recording and a sentence in, the sentence in that voice out, at 24 kHz. /// A module for the reason [`indextts`] is one -- five models, two of them somebody else's. @@ -89,7 +84,9 @@ mod sdxl; /// answer to it there is so far -- a stand-in that makes a noise where the syllables are, so that /// the screen and everything under it can be written before there is a model to put behind it. pub mod speech; -mod suggested; +/// Writing out what a number is read as, so a speech model never sees a digit: the text +/// normalizer IndexTTS-2.5 and CosyVoice3 both read their text through. +pub mod text_normalizer; mod tensor_file; mod tokenizer; /// Reading and writing WAV, which is the one audio format this crate handles itself: a header, a @@ -121,7 +118,6 @@ pub use sdxl::{ PromptEmbedding, SamplerConfig, Sdxl, SdxlConfig, Unet, UnetCondition, UnetConfig, VaeConfig, VaeDecoder, VaeEncoder, VAE_SCALE, }; -pub use suggested::{Size, Suggestions}; pub use tensor_file::read_safetensors; pub use tokenizer::Tokenizer; pub use wav::Sound; diff --git a/waifu/src/manifest.rs b/waifu/src/manifest.rs index 1fbd0ea..583fed0 100644 --- a/waifu/src/manifest.rs +++ b/waifu/src/manifest.rs @@ -52,8 +52,10 @@ //! `weights` names the safetensors files, in the order they are to be read. `tokenizers` names a //! `tokenizer.json` for each way this model turns text into ids -- one for most, two for Anima, //! which reads its prompt through a Qwen3 vocabulary and a T5 one. `config` is what `model.ini` -//! said, block for block and key for key, with nothing renamed on the way across. [`Suggestions`] -//! is what `metadata.json` said, grown from one field to several. +//! said, block for block and key for key, with nothing renamed on the way across. `suggested` is +//! what `metadata.json` said, grown from one field to several: what a screen should start the +//! model's boxes at. Running the model needs none of it, so it is the model manager's to read and +//! this steps over it. //! //! A tokenizer is one line because a tokenizer is one file: everything about how a model's text //! becomes ids -- the vocabulary, the merge ranks, the normalizer, the pre-tokenizer -- is inside @@ -76,7 +78,6 @@ use std::path::{Path, PathBuf}; use crate::error::{Error, Result}; use crate::mapping::Mapping; -use crate::suggested::Suggestions; use crate::yaml::{self, Node}; /// What a model's manifest says. @@ -86,7 +87,6 @@ pub struct Manifest { weights: Vec, tokenizers: BTreeMap, config: BTreeMap, - suggested: Suggestions, explicit: bool, } @@ -98,7 +98,6 @@ impl Manifest { pub const WEIGHTS: &'static str = "weights"; pub const TOKENIZERS: &'static str = "tokenizers"; pub const CONFIG: &'static str = "config"; - pub const SUGGESTED: &'static str = "suggested"; /// Spelled after the tag the model cards carry, which is where the answer comes from -- and /// which is also the spelling the hubs will keep: ModelScope's moderation reverts a commit /// that says the same thing in plainer words, silently, reporting the upload as committed and @@ -170,14 +169,6 @@ impl Manifest { config.insert(name, block); } - // The one block that may be missing altogether. A model with no advice to give has none, - // and that is not an absence worth complaining about: it is what every model had before - // there was anywhere to say it. - let suggested = match top.remove(Self::SUGGESTED) { - Some(node) => Suggestions::from_fields(&node.into_map(Self::SUGGESTED)?), - None => Suggestions::default(), - }; - // What the model draws, which the catalogue also knows for a published one -- this is the // answer for a manifest handed over by path, which is in no catalogue at all. Missing // reads as false: every manifest written before there was anywhere to say this says @@ -187,15 +178,14 @@ impl Manifest { .and_then(|node| node.as_str().map(|said| said.trim() == "true")) .unwrap_or(false); - // A block this build has not heard of is stepped over rather than refused, the way an - // unknown suggestion always was: a manifest from a newer writer still names the weights - // and still says what the model is. + // A block this build has not heard of is stepped over rather than refused -- `suggested` + // among them, which is the model manager's: a manifest from a newer writer still names + // the weights and still says what the model is. Ok(Manifest { path: PathBuf::new(), weights, tokenizers, config, - suggested, explicit, }) } @@ -257,11 +247,6 @@ impl Manifest { self.config.contains_key(name) } - /// What the model suggests being asked for, which for an older model is nothing. - pub fn suggested(&self) -> &Suggestions { - &self.suggested - } - /// Whether what it was trained on means it draws explicit pictures readily, asked for them or /// not. False for a manifest that does not say, which is every manifest written before there /// was anywhere to say it. diff --git a/waifu/src/progress.rs b/waifu/src/progress.rs index 2217822..a3707ca 100644 --- a/waifu/src/progress.rs +++ b/waifu/src/progress.rs @@ -17,7 +17,7 @@ // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -//! How far along a fetch, a read or a run is: what a bar is drawn from, and the words on it. +//! How far along a read or a run is: what a bar is drawn from, and the words on it. use std::time::Instant; @@ -51,22 +51,6 @@ const HELD: f64 = 0.99; const SAYING: f64 = 0.4; const DRAWING: f64 = 0.45; -/// Which package of a model is being fetched, and how far into it. -#[derive(Clone)] -pub struct Fetch { - pub model: String, - /// Which hub it is coming from, once the fetch has settled that. Unknown for the moment - /// before: it is decided by a probe, and a screen that guessed would be naming a hub while - /// the probe is still out. - pub hub: Option<&'static str>, - pub file: String, - pub done: u64, - pub total: Option, - pub part: usize, - /// How many packages there are in all, or zero while that is still unknown. - pub parts: usize, -} - /// A run in flight. #[derive(Clone)] pub struct Run { @@ -150,7 +134,6 @@ impl Convert { pub enum Doing { /// Nothing: either no model has been asked for yet, or the last thing asked for has finished. Nothing, - Fetching(Fetch), /// Reading a model off the disk and onto the device, which is minutes for a large one. Reading { model: String, @@ -168,16 +151,11 @@ impl Doing { /// How far along, between nothing and all of it, or `None` for work with no end in sight. /// - /// A fetch knows how many bytes it is; reading a model does not know anything -- it is one - /// call into the tensor library that returns when it returns -- and a bar drawn for it would - /// be a bar making something up. + /// Reading a model does not know anything -- it is one call into the tensor library that + /// returns when it returns -- and a bar drawn for it would be a bar making something up. pub fn fraction(&self) -> Option { match self { Doing::Nothing | Doing::Reading { .. } => None, - Doing::Fetching(fetch) => fetch - .total - .filter(|total| *total > 0) - .map(|total| fetch.done as f64 / total as f64), Doing::Drawing(run) => Some(drawn(run.progress, run.steps)), Doing::Speaking(say) => Some(spoken(say, Instant::now())), Doing::Converting(convert) => Some(converted(convert)), @@ -188,14 +166,6 @@ impl Doing { pub fn words(&self) -> String { match self { Doing::Nothing => String::new(), - Doing::Fetching(fetch) if fetch.file.is_empty() => format!("fetching {}", fetch.model), - Doing::Fetching(fetch) => match fetch.parts { - 0 => format!("fetching {} -- {}", fetch.model, fetch.file), - parts => format!( - "fetching {} -- {} ({} of {parts})", - fetch.model, fetch.file, fetch.part - ), - }, Doing::Reading { model } => format!("reading {model}"), Doing::Drawing(run) => match run.progress { GenerationProgress::Encoding => "reading the prompt".to_string(), @@ -394,31 +364,6 @@ mod tests { assert!(reading.words().contains("sdxl:base")); } - #[test] - fn a_fetch_of_an_unknown_length_has_no_fraction_either() { - // Which is what a server that did not say how long the file is leaves, and it is also - // what the moment before the first byte looks like. - let mut fetch = Fetch { - model: "sdxl:base".to_string(), - hub: None, - file: "sdxl-base.waifupkg".to_string(), - done: 512, - total: None, - part: 1, - parts: 3, - }; - assert_eq!(Doing::Fetching(fetch.clone()).fraction(), None); - - fetch.total = Some(1024); - assert_eq!(Doing::Fetching(fetch.clone()).fraction(), Some(0.5)); - - // And the words say which model, not only which file: the file names in a package are not - // something anybody recognises. - let words = Doing::Fetching(fetch).words(); - assert!(words.contains("sdxl:base"), "{words}"); - assert!(words.contains("1 of 3"), "{words}"); - } - #[test] fn nothing_happening_is_not_busy() { assert!(!Doing::Nothing.is_busy()); diff --git a/waifu/src/runtime.rs b/waifu/src/runtime.rs index f37600c..d07c505 100644 --- a/waifu/src/runtime.rs +++ b/waifu/src/runtime.rs @@ -17,7 +17,7 @@ // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -//! Where a run goes: which device, where the weights wait, and how to choose for a model. +//! Where a run goes: which device, and where the weights wait. use crate::flint::MemorySnapshot; use crate::{Device, Residency}; @@ -126,7 +126,7 @@ impl DeviceOption { DeviceOption::MetalDiskOffload => Runtime::METAL_DISK_OFFLOAD, DeviceOption::Vulkan => Runtime::on(Device::Vulkan), // Never an offload one: with no model in hand there is nothing to measure against - // the card. [`DeviceOption::resolve_for`] is the answer once there is. + // the card. Which one suits a model is the caller's: see `waifu_devices_json`. DeviceOption::Auto => { // At most one of CUDA and Metal is ever built, so the order between those two only // decides which check runs first. Vulkan comes last because it runs on the same @@ -144,18 +144,6 @@ impl DeviceOption { } } } - - /// The runtime this means for a model of `weights` bytes: a named device is that device, and - /// `auto` is [`Runtime::best`] for this machine as it is now. - pub fn resolve_for(self, weights: Option) -> Runtime { - match self { - DeviceOption::Auto => { - let available = Runtime::available(); - Runtime::best(weights, &available, Room::measure(&available)) - } - named => named.resolve(), - } - } } /// How much of each card is free, where the card said. `None` is a card that was not asked or @@ -189,18 +177,6 @@ impl Room { } } -/// What a model of `weights` bytes wants of a card to run with its weights on it: the weights, -/// and room beside them for what a run makes -- the latents, the attention, the decoded picture. -/// -/// A guess, and a generous one. Too small and the run aborts out of memory halfway through a -/// picture; too large and a model that would have fitted runs offloaded, which is slower but -/// finishes. SDXL's 7 GB comes to 9.7, and Qwen-Image's 15 GB of fp8 to 18.6, which is what puts -/// the one on a sixteen gigabyte card and the other beside it. -pub fn wanted_on_the_card(weights: u64) -> u64 { - const GIB: u64 = 1 << 30; - weights + weights / 10 + 2 * GIB -} - impl Runtime { /// Every place a run can go that this build and this machine can take it. /// @@ -211,135 +187,4 @@ impl Runtime { .filter(|runtime| runtime.device().is_available()) .collect() } - - /// What a model of `weights` bytes wants of this runtime's card and what the card has free, - /// where the one is more than the other: a run there is likely to abort out of memory. - /// - /// Only a runtime that puts the weights on a card that was measured can be short. The offload - /// one keeps them on the host, Metal's memory is the machine's, and a size or a card that is - /// not known is taken as one that fits. - pub fn short_of(self, weights: Option, room: Room) -> Option<(u64, u64)> { - let free = match (self.device, self.residency) { - (Device::Cuda, Residency::Device) => room.cuda, - (Device::Vulkan, Residency::Device) => room.vulkan, - _ => None, - }?; - let wanted = wanted_on_the_card(weights?); - (wanted > free).then_some((wanted, free)) - } - - /// The fastest of `available` a model of `weights` bytes will run on, given `room`. - /// - /// CUDA where the model fits on the card, and CUDA with the weights on the host where it does - /// not: slower, but far faster than the processor, and it finishes where the other aborts. - /// Then Metal, whose memory is the machine's. Then Vulkan, which has no offload of its own and - /// so is only the answer where the model fits. The processor otherwise. A size that is not - /// known is taken as one that fits. - pub fn best(weights: Option, available: &[Runtime], room: Room) -> Runtime { - let cuda = Runtime::on(Device::Cuda); - let metal = Runtime::on(Device::Metal); - let vulkan = Runtime::on(Device::Vulkan); - - if available.contains(&cuda) { - if cuda.short_of(weights, room).is_none() { - return cuda; - } - if available.contains(&Runtime::CUDA_CPU_OFFLOAD) { - return Runtime::CUDA_CPU_OFFLOAD; - } - } - if available.contains(&metal) { - return metal; - } - if available.contains(&vulkan) && vulkan.short_of(weights, room).is_none() { - return vulkan; - } - Runtime::on(Device::Cpu) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - const GB: u64 = 1_000_000_000; - - #[test] - fn auto_puts_a_model_on_the_card_it_fits_and_offloads_one_it_does_not() { - let everything = Runtime::ALL.to_vec(); - let cuda = Runtime::on(Device::Cuda); - let sixteen = Room { - cuda: Some(16 * GB), - vulkan: Some(16 * GB), - }; - - // SDXL on a sixteen gigabyte card, and Qwen-Image's fp8 beside it. - assert_eq!(Runtime::best(Some(7 * GB), &everything, sixteen), cuda); - assert_eq!( - Runtime::best(Some(15 * GB), &everything, sixteen), - Runtime::CUDA_CPU_OFFLOAD - ); - - // A size or a card that did not say is the old answer, which is the card. - assert_eq!(Runtime::best(None, &everything, sixteen), cuda); - assert_eq!( - Runtime::best(Some(40 * GB), &everything, Room::default()), - cuda - ); - } - - #[test] - fn auto_without_cuda_goes_to_vulkan_only_where_the_model_fits() { - let cpu = Runtime::on(Device::Cpu); - let vulkan = Runtime::on(Device::Vulkan); - let here = [cpu, vulkan]; - let eight = Room { - cuda: None, - vulkan: Some(8 * GB), - }; - - assert_eq!(Runtime::best(Some(3 * GB), &here, eight), vulkan); - assert_eq!(Runtime::best(Some(7 * GB), &here, eight), cpu); - - // Metal has the machine's memory, so it has no card to be too big for -- and its offload - // is never chosen for anyone, however large the model: it is asked for by name. - let mac = [cpu, Runtime::on(Device::Metal), Runtime::METAL_DISK_OFFLOAD]; - assert_eq!( - Runtime::best(Some(400 * GB), &mac, Room::default()), - Runtime::on(Device::Metal) - ); - assert_eq!( - Runtime::METAL_DISK_OFFLOAD.short_of(Some(400 * GB), Room::default()), - None - ); - assert_eq!(Runtime::best(Some(1), &[cpu], Room::default()), cpu); - } - - #[test] - fn only_a_card_the_weights_go_on_can_be_short_of_room() { - let eight = Room { - cuda: Some(8 * GB), - vulkan: Some(8 * GB), - }; - let cuda = Runtime::on(Device::Cuda); - - assert_eq!(cuda.short_of(Some(3 * GB), eight), None); - assert_eq!( - cuda.short_of(Some(7 * GB), eight), - Some((wanted_on_the_card(7 * GB), 8 * GB)) - ); - assert!(Runtime::on(Device::Vulkan) - .short_of(Some(7 * GB), eight) - .is_some()); - - // The offload one keeps the weights on the host, and the processor has no card at all. - assert_eq!( - Runtime::CUDA_CPU_OFFLOAD.short_of(Some(7 * GB), eight), - None - ); - assert_eq!(Runtime::on(Device::Cpu).short_of(Some(7 * GB), eight), None); - // Nor is anything short of what was never measured. - assert_eq!(cuda.short_of(None, eight), None); - assert_eq!(cuda.short_of(Some(7 * GB), Room::default()), None); - } } diff --git a/waifu/src/text_normalizer/chinese.rs b/waifu/src/text_normalizer/chinese.rs new file mode 100644 index 0000000..8626997 --- /dev/null +++ b/waifu/src/text_normalizer/chinese.rs @@ -0,0 +1,153 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! Numbers read in Chinese, which groups them by ten thousand rather than by a thousand. + +const DIGITS: [&str; 10] = ["零", "一", "二", "三", "四", "五", "六", "七", "八", "九"]; + +/// 个, 十, 百, 千 -- the places inside one group of four. +const PLACES: [&str; 4] = ["", "十", "百", "千"]; + +/// What every fourth place is worth. Chinese groups by ten thousand, not by a thousand, which +/// is the single thing that makes this different from the other two. +const GROUPS: [&str; 4] = ["", "万", "亿", "万亿"]; + +/// One group of four digits, 1..=9999. +/// +/// `leading` says nothing has been spoken yet, anywhere in the number, which is what decides +/// between 两 and 二. See [`cardinal`]. +/// +/// Two other rules that are easy to state and easy to get wrong. A run of zeros inside a +/// number becomes exactly one 零 however long it is, and a leading 一十 is said 十 -- ten is +/// 十 and not 一十, while a hundred and ten is 一百一十 and not 一百十. +fn group(value: u32, leading: bool) -> String { + let mut out = String::new(); + let mut zero_pending = false; + let mut started = false; + + for place in (0..4).rev() { + let digit = (value / 10u32.pow(place)) % 10; + + if digit == 0 { + if started { + zero_pending = true; + } + continue; + } + + if zero_pending { + out.push_str(DIGITS[0]); + zero_pending = false; + } + + if place == 1 && digit == 1 && !started { + out.push_str(PLACES[1]); + } else { + let two = digit == 2 && place >= 2 && leading && !started; + out.push_str(if two { "两" } else { DIGITS[digit as usize] }); + out.push_str(PLACES[place as usize]); + } + + started = true; + } + + out +} + +/// A whole number, written out. +/// +/// # 两 and 二 are both "two" +/// +/// Which one a 2 takes is decided by where it sits, and the rule is narrower than it first +/// looks: **两 only when the 2 is the first digit spoken in the whole number and a place word +/// follows it.** Everywhere else it is 二. +/// +/// | | | +/// | --- | --- | +/// | 200, 2000, 20000 | 两百, 两千, 两万 -- leading, with a place | +/// | 2, 20 | 二, 二十 -- leading, but the units and the tens take 二 | +/// | 1234, 12000 | 一千**二**百三十四, 一万**二**千 -- not leading | +/// | 2222 | **两**千**二**百**二**十二 -- all four rules in one number | +/// +/// An earlier version of this used the place alone, which reads 1234 as 一千两百三十四. That +/// is wrong and it is wrong quietly: every test built from small numbers passes. +pub fn cardinal(value: u64) -> String { + if value == 0 { + return DIGITS[0].to_string(); + } + + let mut groups = Vec::new(); + let mut rest = value; + while rest > 0 { + groups.push((rest % 10_000) as u32); + rest /= 10_000; + } + + let mut out = String::new(); + for index in (0..groups.len()).rev() { + let value = groups[index]; + + if value == 0 { + // A whole empty group is still one 零, but only if something below it is spoken. + if !out.is_empty() + && groups[..index].iter().any(|group| *group > 0) + && !out.ends_with(DIGITS[0]) + { + out.push_str(DIGITS[0]); + } + continue; + } + + // 一亿零一: a group that does not fill its four places needs a 零 in front of it. + if !out.is_empty() && value < 1000 && !out.ends_with(DIGITS[0]) { + out.push_str(DIGITS[0]); + } + + let leading = out.is_empty(); + + // A group that is exactly 2 and carries 万 or 亿 takes 两 the same way 两百 does -- + // 20000 is 两万. `group` cannot see that, because inside it the 2 is a lone unit. + if leading && value == 2 && index > 0 { + out.push('两'); + } else { + out.push_str(&group(value, leading)); + } + + out.push_str(GROUPS[index]); + } + + out +} + +/// Digit by digit, which is how a year is read: 2026 as 二零二六年. +pub fn digits(text: &str) -> String { + text.chars() + .filter_map(|c| c.to_digit(10)) + .map(|d| DIGITS[d as usize]) + .collect() +} + +/// The part after the point, which is always read digit by digit. +pub fn fraction(text: &str) -> String { + format!("点{}", digits(text)) +} + +pub const NEGATIVE: &str = "负"; +pub const PERCENT_PREFIX: &str = "百分之"; +pub const ORDINAL_PREFIX: &str = "第"; diff --git a/waifu/src/text_normalizer/english.rs b/waifu/src/text_normalizer/english.rs new file mode 100644 index 0000000..de55b1d --- /dev/null +++ b/waifu/src/text_normalizer/english.rs @@ -0,0 +1,150 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! Numbers read in English. + +const ONES: [&str; 20] = [ + "zero", + "one", + "two", + "three", + "four", + "five", + "six", + "seven", + "eight", + "nine", + "ten", + "eleven", + "twelve", + "thirteen", + "fourteen", + "fifteen", + "sixteen", + "seventeen", + "eighteen", + "nineteen", +]; + +const TENS: [&str; 10] = [ + "", "", "twenty", "thirty", "forty", "fifty", "sixty", "seventy", "eighty", "ninety", +]; + +const SCALES: [&str; 5] = ["", "thousand", "million", "billion", "trillion"]; + +fn under_hundred(value: u32) -> String { + if value < 20 { + return ONES[value as usize].to_string(); + } + + let (tens, ones) = (value / 10, value % 10); + if ones == 0 { + TENS[tens as usize].to_string() + } else { + format!("{} {}", TENS[tens as usize], ONES[ones as usize]) + } +} + +fn under_thousand(value: u32) -> String { + if value < 100 { + return under_hundred(value); + } + + let (hundreds, rest) = (value / 100, value % 100); + if rest == 0 { + format!("{} hundred", ONES[hundreds as usize]) + } else { + // "one hundred and one", which is the reading the reference uses too. + format!( + "{} hundred and {}", + ONES[hundreds as usize], + under_hundred(rest) + ) + } +} + +pub fn cardinal(value: u64) -> String { + if value < 20 { + return ONES[value as usize].to_string(); + } + + let mut groups = Vec::new(); + let mut rest = value; + while rest > 0 { + groups.push((rest % 1000) as u32); + rest /= 1000; + } + + let mut parts: Vec = Vec::new(); + for index in (0..groups.len()).rev() { + if groups[index] == 0 { + continue; + } + + let mut part = under_thousand(groups[index]); + if index > 0 { + part.push(' '); + part.push_str(SCALES[index]); + } + parts.push(part); + } + + // "two thousand and twenty six": the last group joins with "and" when it is small enough + // to be a remainder rather than a magnitude of its own. + if parts.len() > 1 && groups[0] > 0 && groups[0] < 100 { + let tail = parts.pop().unwrap(); + format!("{} and {tail}", parts.join(" ")) + } else { + parts.join(" ") + } +} + +/// first, second, twenty first -- the last word takes the suffix and the rest do not. +pub fn ordinal(value: u64) -> String { + let words = cardinal(value); + let cut = words.rfind(' ').map(|at| at + 1).unwrap_or(0); + let (head, last) = words.split_at(cut); + + let changed = match last { + "one" => "first".to_string(), + "two" => "second".to_string(), + "three" => "third".to_string(), + "five" => "fifth".to_string(), + "eight" => "eighth".to_string(), + "nine" => "ninth".to_string(), + "twelve" => "twelfth".to_string(), + other if other.ends_with('y') => format!("{}ieth", &other[..other.len() - 1]), + other => format!("{other}th"), + }; + + format!("{head}{changed}") +} + +pub fn fraction(text: &str) -> String { + let spoken: Vec<&str> = text + .chars() + .filter_map(|c| c.to_digit(10)) + .map(|d| ONES[d as usize]) + .collect(); + + format!("point {}", spoken.join(" ")) +} + +pub const NEGATIVE: &str = "minus"; +pub const PERCENT_SUFFIX: &str = "percent"; diff --git a/waifu/src/indextts/normalize.rs b/waifu/src/text_normalizer/mod.rs similarity index 54% rename from waifu/src/indextts/normalize.rs rename to waifu/src/text_normalizer/mod.rs index 393fa75..12c7710 100644 --- a/waifu/src/indextts/normalize.rs +++ b/waifu/src/text_normalizer/mod.rs @@ -17,7 +17,8 @@ // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -//! Writing out what a number is read as, so a speech model never sees a digit. +//! Writing out what a number is read as, so a speech model never sees a digit: the text normalizer +//! IndexTTS-2.5 and CosyVoice3 both read their text through. //! //! A text-to-speech model is trained on words. `2026` is not a word, and what a model does with //! one is whatever its tokenizer happens to do -- which is why every TTS system has a pass like @@ -28,7 +29,7 @@ //! deliberately: see below. //! //! ``` -//! use waifu::indextts::normalize::{normalize, Language}; +//! use waifu::text_normalizer::{normalize, Language}; //! //! assert_eq!(normalize("共1234人", Language::Chinese), "共一千二百三十四人"); //! assert_eq!(normalize("3.5% more", Language::English), "three point five percent more"); @@ -73,7 +74,7 @@ //! //! Cardinals, decimals, negatives, percentages, money, ordinals, Chinese dates written with //! 年月日, and clock times written `H:MM`. What is deliberately not covered is listed at the -//! bottom of `docs/indextts_normalize.md`, so a caller can see the edge before walking off it. +//! bottom of `docs/text_normalize.md`, so a caller can see the edge before walking off it. /// Which language's rules to read a number by. /// @@ -107,450 +108,10 @@ impl Language { } } -// --------------------------------------------------------------------------------------------- -// Chinese -// --------------------------------------------------------------------------------------------- - -mod chinese { - const DIGITS: [&str; 10] = ["零", "一", "二", "三", "四", "五", "六", "七", "八", "九"]; - - /// 个, 十, 百, 千 -- the places inside one group of four. - const PLACES: [&str; 4] = ["", "十", "百", "千"]; - - /// What every fourth place is worth. Chinese groups by ten thousand, not by a thousand, which - /// is the single thing that makes this different from the other two. - const GROUPS: [&str; 4] = ["", "万", "亿", "万亿"]; - - /// One group of four digits, 1..=9999. - /// - /// `leading` says nothing has been spoken yet, anywhere in the number, which is what decides - /// between 两 and 二. See [`cardinal`]. - /// - /// Two other rules that are easy to state and easy to get wrong. A run of zeros inside a - /// number becomes exactly one 零 however long it is, and a leading 一十 is said 十 -- ten is - /// 十 and not 一十, while a hundred and ten is 一百一十 and not 一百十. - fn group(value: u32, leading: bool) -> String { - let mut out = String::new(); - let mut zero_pending = false; - let mut started = false; - - for place in (0..4).rev() { - let digit = (value / 10u32.pow(place)) % 10; - - if digit == 0 { - if started { - zero_pending = true; - } - continue; - } - - if zero_pending { - out.push_str(DIGITS[0]); - zero_pending = false; - } - - if place == 1 && digit == 1 && !started { - out.push_str(PLACES[1]); - } else { - let two = digit == 2 && place >= 2 && leading && !started; - out.push_str(if two { "两" } else { DIGITS[digit as usize] }); - out.push_str(PLACES[place as usize]); - } - - started = true; - } - - out - } - - /// A whole number, written out. - /// - /// # 两 and 二 are both "two" - /// - /// Which one a 2 takes is decided by where it sits, and the rule is narrower than it first - /// looks: **两 only when the 2 is the first digit spoken in the whole number and a place word - /// follows it.** Everywhere else it is 二. - /// - /// | | | - /// | --- | --- | - /// | 200, 2000, 20000 | 两百, 两千, 两万 -- leading, with a place | - /// | 2, 20 | 二, 二十 -- leading, but the units and the tens take 二 | - /// | 1234, 12000 | 一千**二**百三十四, 一万**二**千 -- not leading | - /// | 2222 | **两**千**二**百**二**十二 -- all four rules in one number | - /// - /// An earlier version of this used the place alone, which reads 1234 as 一千两百三十四. That - /// is wrong and it is wrong quietly: every test built from small numbers passes. - pub fn cardinal(value: u64) -> String { - if value == 0 { - return DIGITS[0].to_string(); - } - - let mut groups = Vec::new(); - let mut rest = value; - while rest > 0 { - groups.push((rest % 10_000) as u32); - rest /= 10_000; - } - - let mut out = String::new(); - for index in (0..groups.len()).rev() { - let value = groups[index]; - - if value == 0 { - // A whole empty group is still one 零, but only if something below it is spoken. - if !out.is_empty() - && groups[..index].iter().any(|group| *group > 0) - && !out.ends_with(DIGITS[0]) - { - out.push_str(DIGITS[0]); - } - continue; - } - - // 一亿零一: a group that does not fill its four places needs a 零 in front of it. - if !out.is_empty() && value < 1000 && !out.ends_with(DIGITS[0]) { - out.push_str(DIGITS[0]); - } - - let leading = out.is_empty(); - - // A group that is exactly 2 and carries 万 or 亿 takes 两 the same way 两百 does -- - // 20000 is 两万. `group` cannot see that, because inside it the 2 is a lone unit. - if leading && value == 2 && index > 0 { - out.push('两'); - } else { - out.push_str(&group(value, leading)); - } - - out.push_str(GROUPS[index]); - } - - out - } - - /// Digit by digit, which is how a year is read: 2026 as 二零二六年. - pub fn digits(text: &str) -> String { - text.chars() - .filter_map(|c| c.to_digit(10)) - .map(|d| DIGITS[d as usize]) - .collect() - } - - /// The part after the point, which is always read digit by digit. - pub fn fraction(text: &str) -> String { - format!("点{}", digits(text)) - } - - pub const NEGATIVE: &str = "负"; - pub const PERCENT_PREFIX: &str = "百分之"; - pub const ORDINAL_PREFIX: &str = "第"; -} - -// --------------------------------------------------------------------------------------------- -// English -// --------------------------------------------------------------------------------------------- - -mod english { - const ONES: [&str; 20] = [ - "zero", - "one", - "two", - "three", - "four", - "five", - "six", - "seven", - "eight", - "nine", - "ten", - "eleven", - "twelve", - "thirteen", - "fourteen", - "fifteen", - "sixteen", - "seventeen", - "eighteen", - "nineteen", - ]; - - const TENS: [&str; 10] = [ - "", "", "twenty", "thirty", "forty", "fifty", "sixty", "seventy", "eighty", "ninety", - ]; - - const SCALES: [&str; 5] = ["", "thousand", "million", "billion", "trillion"]; - - fn under_hundred(value: u32) -> String { - if value < 20 { - return ONES[value as usize].to_string(); - } - - let (tens, ones) = (value / 10, value % 10); - if ones == 0 { - TENS[tens as usize].to_string() - } else { - format!("{} {}", TENS[tens as usize], ONES[ones as usize]) - } - } - - fn under_thousand(value: u32) -> String { - if value < 100 { - return under_hundred(value); - } - - let (hundreds, rest) = (value / 100, value % 100); - if rest == 0 { - format!("{} hundred", ONES[hundreds as usize]) - } else { - // "one hundred and one", which is the reading the reference uses too. - format!( - "{} hundred and {}", - ONES[hundreds as usize], - under_hundred(rest) - ) - } - } - - pub fn cardinal(value: u64) -> String { - if value < 20 { - return ONES[value as usize].to_string(); - } - - let mut groups = Vec::new(); - let mut rest = value; - while rest > 0 { - groups.push((rest % 1000) as u32); - rest /= 1000; - } - - let mut parts: Vec = Vec::new(); - for index in (0..groups.len()).rev() { - if groups[index] == 0 { - continue; - } - - let mut part = under_thousand(groups[index]); - if index > 0 { - part.push(' '); - part.push_str(SCALES[index]); - } - parts.push(part); - } - - // "two thousand and twenty six": the last group joins with "and" when it is small enough - // to be a remainder rather than a magnitude of its own. - if parts.len() > 1 && groups[0] > 0 && groups[0] < 100 { - let tail = parts.pop().unwrap(); - format!("{} and {tail}", parts.join(" ")) - } else { - parts.join(" ") - } - } - - /// first, second, twenty first -- the last word takes the suffix and the rest do not. - pub fn ordinal(value: u64) -> String { - let words = cardinal(value); - let cut = words.rfind(' ').map(|at| at + 1).unwrap_or(0); - let (head, last) = words.split_at(cut); - - let changed = match last { - "one" => "first".to_string(), - "two" => "second".to_string(), - "three" => "third".to_string(), - "five" => "fifth".to_string(), - "eight" => "eighth".to_string(), - "nine" => "ninth".to_string(), - "twelve" => "twelfth".to_string(), - other if other.ends_with('y') => format!("{}ieth", &other[..other.len() - 1]), - other => format!("{other}th"), - }; - - format!("{head}{changed}") - } - - pub fn fraction(text: &str) -> String { - let spoken: Vec<&str> = text - .chars() - .filter_map(|c| c.to_digit(10)) - .map(|d| ONES[d as usize]) - .collect(); - - format!("point {}", spoken.join(" ")) - } - - pub const NEGATIVE: &str = "minus"; - pub const PERCENT_SUFFIX: &str = "percent"; -} - -// --------------------------------------------------------------------------------------------- -// Spanish -// --------------------------------------------------------------------------------------------- - -mod spanish { - const ONES: [&str; 30] = [ - "cero", - "uno", - "dos", - "tres", - "cuatro", - "cinco", - "seis", - "siete", - "ocho", - "nueve", - "diez", - "once", - "doce", - "trece", - "catorce", - "quince", - "dieciséis", - "diecisiete", - "dieciocho", - "diecinueve", - "veinte", - "veintiuno", - "veintidós", - "veintitrés", - "veinticuatro", - "veinticinco", - "veintiséis", - "veintisiete", - "veintiocho", - "veintinueve", - ]; - - const TENS: [&str; 10] = [ - "", - "", - "veinte", - "treinta", - "cuarenta", - "cincuenta", - "sesenta", - "setenta", - "ochenta", - "noventa", - ]; - - /// The hundreds are their own words rather than "n hundred", and 500, 700 and 900 are not - /// formed the way the other six are. - const HUNDREDS: [&str; 10] = [ - "", - "ciento", - "doscientos", - "trescientos", - "cuatrocientos", - "quinientos", - "seiscientos", - "setecientos", - "ochocientos", - "novecientos", - ]; - - fn under_hundred(value: u32) -> String { - if value < 30 { - return ONES[value as usize].to_string(); - } - - let (tens, ones) = (value / 10, value % 10); - if ones == 0 { - TENS[tens as usize].to_string() - } else { - // treinta y uno: from thirty up the two halves are joined by "y". - format!("{} y {}", TENS[tens as usize], ONES[ones as usize]) - } - } - - fn under_thousand(value: u32) -> String { - if value < 100 { - return under_hundred(value); - } - - let (hundreds, rest) = (value / 100, value % 100); - if rest == 0 { - // Exactly one hundred is "cien"; a hundred and something is "ciento ...". - return if hundreds == 1 { - "cien".to_string() - } else { - HUNDREDS[hundreds as usize].to_string() - }; - } - - format!("{} {}", HUNDREDS[hundreds as usize], under_hundred(rest)) - } - - pub fn cardinal(value: u64) -> String { - if value < 30 { - return ONES[value as usize].to_string(); - } - - let millions = value / 1_000_000; - let rest = value % 1_000_000; - - let mut parts: Vec = Vec::new(); - - if millions > 0 { - // "un millón", not "uno millón", and the plural changes the noun rather than a suffix. - parts.push(if millions == 1 { - "un millón".to_string() - } else { - format!("{} millones", cardinal(millions)) - }); - } - - let thousands = rest / 1000; - if thousands > 0 { - parts.push(if thousands == 1 { - "mil".to_string() - } else { - format!("{} mil", under_thousand(thousands as u32)) - }); - } - - let units = rest % 1000; - if units > 0 { - parts.push(under_thousand(units as u32)); - } - - if parts.is_empty() { - under_thousand(0) - } else { - parts.join(" ") - } - } - - /// primero through décimo, and the cardinal beyond -- which is what Spanish itself tends to - /// do past ten in speech. - pub fn ordinal(value: u64) -> String { - match value { - 1 => "primero".to_string(), - 2 => "segundo".to_string(), - 3 => "tercero".to_string(), - 4 => "cuarto".to_string(), - 5 => "quinto".to_string(), - 6 => "sexto".to_string(), - 7 => "séptimo".to_string(), - 8 => "octavo".to_string(), - 9 => "noveno".to_string(), - 10 => "décimo".to_string(), - other => cardinal(other), - } - } - - pub fn fraction(text: &str) -> String { - let spoken: Vec<&str> = text - .chars() - .filter_map(|c| c.to_digit(10)) - .map(|d| ONES[d as usize]) - .collect(); - - format!("coma {}", spoken.join(" ")) - } - - pub const NEGATIVE: &str = "menos"; - pub const PERCENT_SUFFIX: &str = "por ciento"; -} +/// The rules for each language, one file each. +mod chinese; +mod english; +mod spanish; // --------------------------------------------------------------------------------------------- // What a currency symbol is called diff --git a/waifu/src/text_normalizer/spanish.rs b/waifu/src/text_normalizer/spanish.rs new file mode 100644 index 0000000..b7c65a5 --- /dev/null +++ b/waifu/src/text_normalizer/spanish.rs @@ -0,0 +1,184 @@ +// The MIT License (MIT) +// +// Copyright (c) 2026 Xiaoyang Chen +// +// Permission is hereby granted, free of charge, to any person obtaining a copy of this software +// and associated documentation files (the "Software"), to deal in the Software without +// restriction, including without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +// Software is furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! Numbers read in Spanish. + +const ONES: [&str; 30] = [ + "cero", + "uno", + "dos", + "tres", + "cuatro", + "cinco", + "seis", + "siete", + "ocho", + "nueve", + "diez", + "once", + "doce", + "trece", + "catorce", + "quince", + "dieciséis", + "diecisiete", + "dieciocho", + "diecinueve", + "veinte", + "veintiuno", + "veintidós", + "veintitrés", + "veinticuatro", + "veinticinco", + "veintiséis", + "veintisiete", + "veintiocho", + "veintinueve", +]; + +const TENS: [&str; 10] = [ + "", + "", + "veinte", + "treinta", + "cuarenta", + "cincuenta", + "sesenta", + "setenta", + "ochenta", + "noventa", +]; + +/// The hundreds are their own words rather than "n hundred", and 500, 700 and 900 are not +/// formed the way the other six are. +const HUNDREDS: [&str; 10] = [ + "", + "ciento", + "doscientos", + "trescientos", + "cuatrocientos", + "quinientos", + "seiscientos", + "setecientos", + "ochocientos", + "novecientos", +]; + +fn under_hundred(value: u32) -> String { + if value < 30 { + return ONES[value as usize].to_string(); + } + + let (tens, ones) = (value / 10, value % 10); + if ones == 0 { + TENS[tens as usize].to_string() + } else { + // treinta y uno: from thirty up the two halves are joined by "y". + format!("{} y {}", TENS[tens as usize], ONES[ones as usize]) + } +} + +fn under_thousand(value: u32) -> String { + if value < 100 { + return under_hundred(value); + } + + let (hundreds, rest) = (value / 100, value % 100); + if rest == 0 { + // Exactly one hundred is "cien"; a hundred and something is "ciento ...". + return if hundreds == 1 { + "cien".to_string() + } else { + HUNDREDS[hundreds as usize].to_string() + }; + } + + format!("{} {}", HUNDREDS[hundreds as usize], under_hundred(rest)) +} + +pub fn cardinal(value: u64) -> String { + if value < 30 { + return ONES[value as usize].to_string(); + } + + let millions = value / 1_000_000; + let rest = value % 1_000_000; + + let mut parts: Vec = Vec::new(); + + if millions > 0 { + // "un millón", not "uno millón", and the plural changes the noun rather than a suffix. + parts.push(if millions == 1 { + "un millón".to_string() + } else { + format!("{} millones", cardinal(millions)) + }); + } + + let thousands = rest / 1000; + if thousands > 0 { + parts.push(if thousands == 1 { + "mil".to_string() + } else { + format!("{} mil", under_thousand(thousands as u32)) + }); + } + + let units = rest % 1000; + if units > 0 { + parts.push(under_thousand(units as u32)); + } + + if parts.is_empty() { + under_thousand(0) + } else { + parts.join(" ") + } +} + +/// primero through décimo, and the cardinal beyond -- which is what Spanish itself tends to +/// do past ten in speech. +pub fn ordinal(value: u64) -> String { + match value { + 1 => "primero".to_string(), + 2 => "segundo".to_string(), + 3 => "tercero".to_string(), + 4 => "cuarto".to_string(), + 5 => "quinto".to_string(), + 6 => "sexto".to_string(), + 7 => "séptimo".to_string(), + 8 => "octavo".to_string(), + 9 => "noveno".to_string(), + 10 => "décimo".to_string(), + other => cardinal(other), + } +} + +pub fn fraction(text: &str) -> String { + let spoken: Vec<&str> = text + .chars() + .filter_map(|c| c.to_digit(10)) + .map(|d| ONES[d as usize]) + .collect(); + + format!("coma {}", spoken.join(" ")) +} + +pub const NEGATIVE: &str = "menos"; +pub const PERCENT_SUFFIX: &str = "por ciento"; diff --git a/waifu/src/yaml.rs b/waifu/src/yaml.rs index c26bcf0..c3b3ec8 100644 --- a/waifu/src/yaml.rs +++ b/waifu/src/yaml.rs @@ -71,13 +71,6 @@ impl Node { } } - /// The items of this value, if it is a list. - pub fn as_seq(&self) -> Option<&[Node]> { - match self { - Node::Seq(items) => Some(items), - _ => None, - } - } } /// One line that is neither blank nor a comment: how far it is indented, what is on it, and which diff --git a/waifu/tests/indextts_emotion.rs b/waifu/tests/indextts_emotion.rs index 8e37bce..513f8c6 100644 --- a/waifu/tests/indextts_emotion.rs +++ b/waifu/tests/indextts_emotion.rs @@ -23,8 +23,8 @@ //! Every other reference here sweeps a sine per parameter. This draws by xorshift instead, //! because most of what this module does is attention: a swept parameter makes every key in a //! projection point nearly the same way, the softmax over them comes out flat, and a flat softmax -//! is an average -- which does not care what order its keys arrived in. `docs/indextts_gpt.md` -//! has the longer version, and the bug it cost. +//! is an average -- which does not care what order its keys arrived in. `docs/indextts.md` +//! says why, under the GPT. //! //! # Four probes and a table //! diff --git a/waifu/tests/manifest.rs b/waifu/tests/manifest.rs index 44420d3..b752fce 100644 --- a/waifu/tests/manifest.rs +++ b/waifu/tests/manifest.rs @@ -69,9 +69,9 @@ fn temp_dir(what: &str) -> std::path::PathBuf { dir } -/// A whole manifest whose one suggestion is `line`, written as it stands. -fn suggesting(line: &str) -> String { - format!("weights:\n - m.safetensors\nconfig:\n model:\n type: sdxl\nsuggested:\n {line}\n") +/// A whole manifest whose model block holds `line`, written as it stands, besides its type. +fn configuring(line: &str) -> String { + format!("weights:\n - m.safetensors\nconfig:\n model:\n type: sdxl\n {line}\n") } #[test] @@ -127,16 +127,8 @@ fn reads_the_three_blocks_a_manifest_merged() { assert_eq!(sdxl.get_or("force_upcast", 1).unwrap(), 1); assert_eq!(sdxl.get::("clip_skip").unwrap(), 2); - // What used to be metadata.json, which has grown from one field to several. - let suggested = manifest.suggested(); - assert_eq!( - suggested.prompt.as_deref(), - Some("masterpiece, best quality") - ); - assert_eq!(suggested.sizes, vec![(1024, 1024), (832, 1216), (1216, 832)]); - assert_eq!(suggested.size(), Some((1024, 1024))); - assert_eq!(suggested.steps, Some(8)); - assert_eq!(suggested.guidance, Some(1.0)); + // What used to be metadata.json is the model manager's, and stepped over here: reading the + // manifest above at all is the check that it is. assert!(manifest.has_section("sdxl")); assert!(!manifest.has_section("anima")); @@ -181,18 +173,6 @@ fn a_model_that_tokenizes_nothing_names_no_tokenizers() { assert!(manifest.tokenizer("tokenizer").is_err()); } -#[test] -fn a_model_with_no_advice_to_give_has_no_suggested_block() { - let manifest = - Manifest::parse("weights:\n - m.safetensors\nconfig:\n model:\n type: sdxl\n").unwrap(); - - let suggested = manifest.suggested(); - assert_eq!(suggested.prompt, None); - assert_eq!(suggested.steps, None); - assert_eq!(suggested.guidance, None); - assert!(suggested.sizes.is_empty()); -} - #[test] fn a_block_this_build_has_not_heard_of_is_stepped_over() { // A manifest from a newer writer still names the weights and still says what the model is. @@ -209,60 +189,48 @@ fn a_block_this_build_has_not_heard_of_is_stepped_over() { manifest.section("model").unwrap().get_str("type").unwrap(), "sdxl" ); - assert_eq!(manifest.suggested().prompt.as_deref(), Some("1girl")); } #[test] fn scalars_mean_what_they_mean_everywhere_else() { - let prompt = |line: &str| { - Manifest::parse(&suggesting(line)) + let value = |line: &str| { + Manifest::parse(&configuring(line)) + .unwrap() + .section("model") + .unwrap() + .get_str("note") .unwrap() - .suggested() - .prompt - .clone() + .to_string() }; // A space and then a `#` begins a comment, the way it does everywhere else; a `#` with no - // space before it is part of the value. A prompt that really wants the first has to be + // space before it is part of the value. A value that really wants the first has to be // quoted, and the exporter quotes it. - assert_eq!( - prompt("prompt: a plain value # and a comment").as_deref(), - Some("a plain value") - ); - assert_eq!( - prompt("prompt: tag#1 keeps its hash").as_deref(), - Some("tag#1 keeps its hash") - ); - assert_eq!( - prompt(r#"prompt: "tag #1, quoted""#).as_deref(), - Some("tag #1, quoted") - ); + assert_eq!(value("note: a plain value # and a comment"), "a plain value"); + assert_eq!(value("note: tag#1 keeps its hash"), "tag#1 keeps its hash"); + assert_eq!(value(r#"note: "tag #1, quoted""#), "tag #1, quoted"); // Quoted, with the escapes each kind of quote has and no others. - assert_eq!( - prompt(r#"prompt: "a: colon, and a \"quote\"""#).as_deref(), - Some(r#"a: colon, and a "quote""#) - ); - assert_eq!(prompt("prompt: 'it''s here'").as_deref(), Some("it's here")); + assert_eq!(value(r#"note: "a: colon, and a \"quote\"""#), r#"a: colon, and a "quote""#); + assert_eq!(value("note: 'it''s here'"), "it's here"); // Anything wider than ASCII arrives as itself: the exporter writes utf-8 rather than escapes. - assert_eq!(prompt("prompt: 猫, 🐱").as_deref(), Some("猫, 🐱")); + assert_eq!(value("note: 猫, 🐱"), "猫, 🐱"); } #[test] fn a_list_may_be_written_on_one_line_or_on_several() { - let block = Manifest::parse(&suggesting("sizes:\n - [1024, 1024]\n - [832, 1216]")) + let header = "config:\n model:\n type: sdxl\n"; + let block = Manifest::parse(&format!("weights:\n - a.safetensors\n - b.safetensors\n{header}")) .unwrap() - .suggested() - .sizes - .clone(); - let inline = Manifest::parse(&suggesting("sizes: [[1024, 1024], [832, 1216]]")) + .weights() + .to_vec(); + let inline = Manifest::parse(&format!("weights: [a.safetensors, b.safetensors]\n{header}")) .unwrap() - .suggested() - .sizes - .clone(); + .weights() + .to_vec(); - assert_eq!(block, vec![(1024, 1024), (832, 1216)]); + assert_eq!(block, ["a.safetensors", "b.safetensors"]); assert_eq!(block, inline); } diff --git a/waifu/tests/normalize.rs b/waifu/tests/normalize.rs index bde8c02..187074b 100644 --- a/waifu/tests/normalize.rs +++ b/waifu/tests/normalize.rs @@ -9,7 +9,7 @@ //! where this and WeTextProcessing differ, so every divergence is one somebody chose. The cases //! marked "agrees with the reference" are the ones that came out of running it. -use waifu::indextts::normalize::{normalize, Language}; +use waifu::text_normalizer::{normalize, Language}; const ZH: Language = Language::Chinese; const EN: Language = Language::English;