diff --git a/CMakeLists.txt b/CMakeLists.txt index 5fd53e10..891b3c8d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -532,6 +532,7 @@ add_library(engine_core OBJECT src/framework/codecs/mimi_codec_runtime.cpp src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp src/framework/codecs/neural_audio.cpp + src/framework/codecs/oobleck_audio_vae_runtime.cpp src/framework/codecs/redae_codec_runtime.cpp src/framework/conditioners/clap_audio_conditioner_runtime.cpp src/framework/conditioners/cav_mae_st_conditioner_runtime.cpp @@ -1413,6 +1414,38 @@ audiocpp_add_model(stable_audio engine::models::stable_audio::make_stable_audio_loader ) +audiocpp_add_model(yue2 + SOURCES + src/models/yue2/ar_runtime.cpp + src/models/yue2/assets.cpp + src/models/yue2/nar_runtime.cpp + src/models/yue2/pipeline.cpp + src/models/yue2/request.cpp + src/models/yue2/session.cpp + src/models/yue2/tokenizer_text.cpp + src/models/yue2/types.cpp + INCLUDES + engine/models/yue2/session.h + LOADERS + engine::models::yue2::make_yue2_loader +) + +audiocpp_add_model(sheetsage2 + SOURCES + src/models/sheetsage/audio_frontend.cpp + src/models/sheetsage/processing.cpp + src/models/sheetsage/runtime.cpp + src/models/sheetsage/session.cpp + INCLUDES + engine/models/sheetsage/audio_frontend.h + engine/models/sheetsage/processing.h + engine/models/sheetsage/types.h + engine/models/sheetsage/runtime.h + engine/models/sheetsage/session.h + LOADERS + engine::models::sheetsage::make_sheetsage2_loader +) + audiocpp_add_model(supertonic SOURCES src/models/supertonic/assets.cpp @@ -2329,6 +2362,25 @@ if (ENGINE_ENABLE_OPENMP) target_link_libraries(model_perf PRIVATE OpenMP::OpenMP_CXX) endif() +if (ENGINE_BUILD_TESTS OR ENGINE_BUILD_EXTENDED_TESTS OR ENGINE_BUILD_MODEL_TESTS) + add_executable(yue2_vae_parity_probe + tests/yue2/yue2_vae_parity_probe.cpp + ) + target_link_libraries(yue2_vae_parity_probe PRIVATE engine_runtime ggml) + if (ENGINE_ENABLE_OPENMP) + target_link_libraries(yue2_vae_parity_probe PRIVATE OpenMP::OpenMP_CXX) + endif() + + add_executable(sheetsage2_decoder_parity_probe + src/models/sheetsage/runtime.cpp + tests/yue2/sheetsage2_decoder_parity_probe.cpp + ) + target_link_libraries(sheetsage2_decoder_parity_probe PRIVATE engine_runtime ggml) + if (ENGINE_ENABLE_OPENMP) + target_link_libraries(sheetsage2_decoder_parity_probe PRIVATE OpenMP::OpenMP_CXX) + endif() +endif() + if (ENGINE_BUILD_WARMBENCH) function(add_engine_warmbench target_name source_file) add_executable(${target_name} @@ -2468,6 +2520,13 @@ if (ENGINE_BUILD_TESTS OR ENGINE_BUILD_EXTENDED_TESTS OR ENGINE_BUILD_MODEL_TEST add_engine_unittest(audio_dsp_test tests/unittests/test_audio_dsp.cpp) add_test(NAME audio_dsp_test COMMAND audio_dsp_test) + if(sheetsage2 IN_LIST AUDIOCPP_LINKED_MODELS) + add_engine_unittest(sheetsage_audio_frontend_test tests/unittests/test_sheetsage_audio_frontend.cpp) + add_test(NAME sheetsage_audio_frontend_test COMMAND sheetsage_audio_frontend_test --log) + add_engine_unittest(sheetsage_processing_test tests/unittests/test_sheetsage_processing.cpp) + add_test(NAME sheetsage_processing_test COMMAND sheetsage_processing_test --log) + endif() + add_engine_unittest(midi_file_test tests/unittests/test_midi_file.cpp) add_test(NAME midi_file_test COMMAND midi_file_test) diff --git a/docs/models/yue2.md b/docs/models/yue2.md new file mode 100644 index 00000000..bd1d6426 --- /dev/null +++ b/docs/models/yue2.md @@ -0,0 +1,151 @@ +# YuE2 + +YuE2 is wired as `--family yue2 --task gen`. It generates music from lyrics and +a style prompt, with optional symbolic ABC conditioning. + +## Quick Start + +Default packaged GGUF layout: + +```bash +./build/debug/bin/audiocpp_cli \ + --task gen \ + --family yue2 \ + --model models/Yue2-3B-GGUF \ + --backend cuda \ + --threads 8 \ + --lyrics "[Verse] +Soft morning light is touching the window. +I hear the city waking below. +[Chorus] +Stay with the rhythm, let it carry us home. +Sing with the sunrise, we are never alone." \ + --request-option style="English, indie pop, bright acoustic guitar, soft drums, warm lead vocal, polished demo mix" \ + --request-option cot=off \ + --seed 831001 \ + --out yue2.wav \ + --log +``` + +The default session loads `yue2-3b-q8_0.gguf` for the main AR/NAR model and +`yue2-vae-f16.gguf` for the VAE from the model root. + +## Model + +| Field | Value | +|---|---| +| Family | `yue2` | +| Model directory | `models/Yue2-3B-GGUF` | +| Task | `gen` | +| Main GGUF default | `yue2-3b-q8_0.gguf` | +| VAE GGUF default | `yue2-vae-f16.gguf` | +| Required sidecars | `sidecars/yue2-model-config.json`, `sidecars/yue2-generation-config.json`, `sidecars/yue2-qwen.tiktoken`, `sidecars/yue2-vae-config.json` | +| Lyrics input | `--lyrics`; `--text` is accepted as a fallback | +| Style input | `--request-option style=` | + +## Component Selection + +Select the BF16 main model: + +```bash +./build/debug/bin/audiocpp_cli \ + --task gen \ + --family yue2 \ + --model models/Yue2-3B-GGUF \ + --backend cuda \ + --threads 8 \ + --session-option yue2.model_gguf=yue2-3b-bf16.gguf \ + --session-option yue2.vae_gguf=yue2-vae-f16.gguf \ + --lyrics "[Verse] +Soft morning light is touching the window. +[Chorus] +Stay with the rhythm, let it carry us home." \ + --request-option style="English, pop rock, bright guitars, clean drums, warm vocal" \ + --request-option cot=off \ + --seed 831001 \ + --out yue2-bf16.wav \ + --log +``` + +Select the F32 VAE: + +```bash +--session-option yue2.vae_gguf=yue2-vae-f32.gguf +``` + +The component paths are relative to `--model`; absolute paths are rejected. + +## ABC Conditioning + +Use `cot=melody` or `cot=full` to run the symbolic route. External ABC requires +one of those modes: + +```bash +./build/debug/bin/audiocpp_cli \ + --task gen \ + --family yue2 \ + --model models/Yue2-3B-GGUF \ + --backend cuda \ + --threads 8 \ + --lyrics "[Verse] +Write the melody over this score." \ + --request-option style="English, folk pop, acoustic guitar, steady drums" \ + --request-option cot=melody \ + --request-option abc_file=/path/to/score.abc \ + --seed 831001 \ + --out yue2-abc.wav \ + --log +``` + +Inline ABC can be passed with `--request-option abc=`. + +## Request Options + +| Option | Values | Default | Meaning | +|---|---|---:|---| +| `--lyrics` | text | required | Song lyrics. | +| `--text` | text | empty | Fallback lyrics source when `--lyrics` is not supplied. | +| `--request-option style=` | text | required | Music style prompt. | +| `--request-option cot=` | `off`, `melody`, `full` | `full` | Symbolic planning route. | +| `--request-option abc=` | ABC text | empty | Inline ABC score; requires `cot=melody` or `cot=full`. | +| `--request-option abc_file=` | path | empty | ABC score file; requires `cot=melody` or `cot=full`. | +| `--request-option semantic_codes_file=` | raw int32 file | empty | Teacher-forced semantic codec IDs for parity/debug runs. | +| `--request-option nar_noise_file=` | raw float32 file | empty | Teacher-forced NAR noise rows with 64 columns for parity/debug runs. | +| `--request-option cfg_scale=` | `0..20` | `1.01` for `cot=off`, otherwise `1.0` | Semantic classifier-free guidance scale. | +| `--request-option num_inference_steps=` | integer > 0 | `32` | NAR midpoint ODE steps. | +| `--seed ` | integer in `[0, 2^63)` | `831001` | Generation seed. Equivalent to `--request-option seed=`. | + +## Sampling Options + +| Option | Values | Default | Meaning | +|---|---|---:|---| +| `--request-option abc_temperature=` | `0..5` | `0.7` | ABC planner sampling temperature. | +| `--request-option abc_top_p=` | `0..1` | `0.9` | ABC planner nucleus sampling probability. | +| `--request-option abc_top_k=` | integer >= 1 | `30` | ABC planner top-k limit. | +| `--request-option abc_repetition_penalty=` | float > 0 | `1.005` | ABC planner repetition penalty. | +| `--request-option abc_penalty_window=` | integer >= 1 | `100` | ABC planner repetition penalty window. | +| `--request-option abc_min_tokens=` | integer >= 0 | `32` | Minimum ABC planner tokens before EOS is accepted. | +| `--request-option abc_max_tokens=` | integer >= `abc_min_tokens` | `4096` | Maximum ABC planner tokens. | +| `--request-option semantic_temperature=` | `0..5` | `1.0` | Semantic codec sampling temperature. | +| `--request-option semantic_top_p=` | `0..1` | `0.95` | Semantic codec nucleus sampling probability. | +| `--request-option semantic_top_k=` | integer >= 1 | `100` | Semantic codec top-k limit. | +| `--request-option semantic_repetition_penalty=` | float > 0 | `1.2` | Semantic codec repetition penalty. | +| `--request-option semantic_penalty_window=` | integer >= 1 | `50` | Semantic codec repetition penalty window. | +| `--request-option semantic_min_tokens=` | integer >= 0 | `200` | Minimum semantic tokens before EOS is accepted. | +| `--request-option semantic_max_tokens=` | integer >= `semantic_min_tokens` | `9000` | Maximum semantic codec tokens. | + +## Session Options + +| Option | Values | Default | Meaning | +|---|---|---:|---| +| `--session-option yue2.model_gguf=` | relative GGUF path | `yue2-3b-q8_0.gguf` | Main AR/NAR component. | +| `--session-option yue2.vae_gguf=` | relative GGUF path | `yue2-vae-f16.gguf` | VAE component. | +| `--session-option yue2.weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0`, `q4_0`, `q4_k` | `native` | Shared weight storage fallback for the main model and VAE. | +| `--session-option yue2.model_weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0`, `q4_0`, `q4_k` | `native` | Main model weight storage override. | +| `--session-option yue2.vae_weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0`, `q4_0`, `q4_k` | `native` | VAE weight storage override. | +| `--session-option yue2.model_weight_context_mb=` | MiB integer >= 1 | `6144` | Main model weight context size. | +| `--session-option yue2.vae_weight_context_mb=` | MiB integer >= 1 | `1536` | VAE weight context size. | +| `--session-option yue2.ar_prefill_graph_arena_mb=` | MiB integer >= 1 | `4096` | AR prefill graph arena size. | +| `--session-option yue2.ar_decode_graph_arena_mb=` | MiB integer >= 1 | `1536` | AR one-token decode graph arena size. | +| `--session-option yue2.nar_graph_arena_mb=` | MiB integer >= 1 | `6144` | NAR acoustic flow graph arena size. | +| `--session-option yue2.vae_graph_arena_mb=` | MiB integer >= 1 | `1536` | VAE decode graph arena size. | diff --git a/include/engine/framework/audio/conversion.h b/include/engine/framework/audio/conversion.h index 906e0127..79c9d59d 100644 --- a/include/engine/framework/audio/conversion.h +++ b/include/engine/framework/audio/conversion.h @@ -1,5 +1,6 @@ #pragma once +#include "engine/framework/audio/resampling.h" #include "engine/framework/audio/wav_reader.h" #include @@ -46,6 +47,14 @@ std::vector convert_interleaved_audio_to_mono_linear_resampled( int channel_count, int target_sample_rate_hz); +std::vector convert_interleaved_audio_to_mono_torchaudio_sinc_hann_resampled( + const std::vector & interleaved_samples, + int sample_rate_hz, + int channel_count, + int target_sample_rate_hz, + const TorchaudioSincHannResampleOptions & options = {}, + MonoMixAccumulation accumulation = MonoMixAccumulation::Float32); + std::vector read_wav_f32_as_mono_linear_resampled( const std::filesystem::path & path, int target_sample_rate_hz); diff --git a/include/engine/framework/codecs/oobleck_audio_vae_runtime.h b/include/engine/framework/codecs/oobleck_audio_vae_runtime.h new file mode 100644 index 00000000..d603e9c4 --- /dev/null +++ b/include/engine/framework/codecs/oobleck_audio_vae_runtime.h @@ -0,0 +1,65 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include +#include + +namespace engine::codecs { + +struct OobleckAudioVaeConfig { + int64_t sample_rate = 48000; + int64_t audio_channels = 2; + int64_t channels = 64; + int64_t encoder_latent_dim = 128; + int64_t decoder_latent_dim = 64; + std::vector c_mults{1, 2, 4, 8, 16, 32}; + std::vector strides{2, 2, 4, 4, 5, 6}; + bool use_snake = true; + bool snake_logscale = true; + bool final_tanh = false; + std::string encoder_prefix = "encoder"; + std::string decoder_prefix = "decoder"; +}; + +struct OobleckAudioVaeRuntimeOptions { + size_t graph_arena_bytes = 512ull * 1024ull * 1024ull; + size_t weight_context_bytes = 1400ull * 1024ull * 1024ull; + assets::TensorStorageType weight_storage_type = assets::TensorStorageType::Native; +}; + +class OobleckAudioVaeRuntime { +public: + OobleckAudioVaeRuntime( + std::shared_ptr source, + core::ExecutionContext & execution, + OobleckAudioVaeConfig config = {}, + OobleckAudioVaeRuntimeOptions options = {}); + ~OobleckAudioVaeRuntime(); + + OobleckAudioVaeRuntime(const OobleckAudioVaeRuntime &) = delete; + OobleckAudioVaeRuntime & operator=(const OobleckAudioVaeRuntime &) = delete; + OobleckAudioVaeRuntime(OobleckAudioVaeRuntime &&) noexcept; + OobleckAudioVaeRuntime & operator=(OobleckAudioVaeRuntime &&) noexcept; + + std::vector encode_planar(const std::vector & planar_audio, int64_t frames); + std::vector decode_planar(const std::vector & latents, int64_t batch, int64_t latent_frames); + std::vector decode(const std::vector & latents, int64_t batch, int64_t latent_frames); + + void prepare_encode(int64_t frames); + void prepare_decode(int64_t batch, int64_t latent_frames); + void release_runtime_graphs(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::codecs diff --git a/include/engine/framework/modules/activation_modules.h b/include/engine/framework/modules/activation_modules.h index 315c6897..b999650b 100644 --- a/include/engine/framework/modules/activation_modules.h +++ b/include/engine/framework/modules/activation_modules.h @@ -154,6 +154,32 @@ class Snake1dModule { Snake1dConfig config_; }; +struct SnakeBeta1dConfig { + int64_t hidden_size = 0; + bool logscale = true; +}; + +struct SnakeBeta1dWeights { + core::TensorValue alpha; + core::TensorValue beta; +}; + +class SnakeBeta1dModule { +public: + explicit SnakeBeta1dModule(SnakeBeta1dConfig config); + + const SnakeBeta1dConfig & config() const noexcept; + const core::ModuleSchema & schema() const noexcept; + core::TensorValue build( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const SnakeBeta1dWeights & weights) const; + static const core::ModuleSchema & static_schema() noexcept; + +private: + SnakeBeta1dConfig config_; +}; + enum class AliasFreeActivationKind { SnakeBeta, }; diff --git a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h index 37c2f04f..e48ce6be 100644 --- a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h +++ b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h @@ -44,6 +44,13 @@ struct QwenCausalPrefillResult { runtime::TransformerKVState state; }; +struct QwenCausalPrefillIntoDecodeResult { + std::vector logits; + std::vector hidden; + int64_t current_end = 0; + int64_t valid_steps = 0; +}; + struct QwenCausalBatchedPrefillResult { std::vector logits; std::vector hidden; @@ -68,6 +75,9 @@ class QwenCausalDecodeRuntime { QwenCausalPrefillResult prefill_tokens(const std::vector & token_ids); QwenCausalPrefillResult prefill_embeddings(const std::vector & embeddings, int64_t steps); + QwenCausalPrefillIntoDecodeResult prefill_tokens_into_decode_cache( + const std::vector & token_ids, + int64_t required_cache_steps); // Prefill bounded blocks directly into the token-decode cache on the backend. // No host KV export/import; subsequent decode_token calls continue this state. @@ -86,6 +96,7 @@ class QwenCausalDecodeRuntime { void start_decode_tokens(const runtime::TransformerKVState & state, int64_t required_cache_steps); void start_decode_embeddings(const runtime::TransformerKVState & state, int64_t required_cache_steps); QwenCausalDecodeStepResult decode_token(int32_t token); + void decode_token_into(int32_t token, QwenCausalDecodeStepResult & out); QwenCausalDecodeStepResult decode_embedding(const std::vector & embedding); void start_decode_tokens_batched( diff --git a/include/engine/framework/runtime/kv_cache.h b/include/engine/framework/runtime/kv_cache.h index f7775916..e8671e36 100644 --- a/include/engine/framework/runtime/kv_cache.h +++ b/include/engine/framework/runtime/kv_cache.h @@ -97,6 +97,8 @@ struct BatchedKVLayerState { struct TransformerBatchedKVState { int64_t batch_size = 0; int64_t current_end = 0; + std::vector current_end_by_batch; + std::vector valid_steps_by_batch; std::vector layers; }; @@ -126,6 +128,8 @@ class TransformerBatchedKVCache { int64_t valid_steps() const noexcept; int64_t current_end() const noexcept; int64_t cache_steps() const noexcept; + const std::vector & valid_steps_by_batch() const noexcept; + const std::vector & current_end_by_batch() const noexcept; private: struct LayerCache { @@ -140,6 +144,8 @@ class TransformerBatchedKVCache { int64_t row_elems_ = 0; int64_t valid_steps_ = 0; int64_t current_end_ = 0; + std::vector valid_steps_by_batch_; + std::vector current_end_by_batch_; TransformerKVCacheOptions options_; std::vector layers_; }; diff --git a/include/engine/models/sheetsage/audio_frontend.h b/include/engine/models/sheetsage/audio_frontend.h new file mode 100644 index 00000000..4a5558c1 --- /dev/null +++ b/include/engine/models/sheetsage/audio_frontend.h @@ -0,0 +1,22 @@ +#pragma once + +#include +#include +#include + +namespace engine::models::sheetsage { + +class SheetSage2AudioFrontend { +public: + std::vector prepare( + const std::vector & interleaved, + int source_rate, + int channels, + int target_rate, + int threads); + +private: + std::map, std::vector> filters_; +}; + +} // namespace engine::models::sheetsage diff --git a/include/engine/models/sheetsage/processing.h b/include/engine/models/sheetsage/processing.h new file mode 100644 index 00000000..83414141 --- /dev/null +++ b/include/engine/models/sheetsage/processing.h @@ -0,0 +1,51 @@ +#pragma once + +#include +#include +#include +#include +#include + +namespace engine::models::sheetsage { + +struct SheetSage2Note { + int pitch = 0; + int track = 0; + int duration_bin = 0; + int duration_steps = 0; +}; + +struct SheetSage2Event { + int64_t subbeat = 0; + int64_t source_subbeat = 0; + int64_t global_subbeat = 0; + int window_index = 0; + float window_start = 0.0F; + float time = 0.0F; + std::optional timestamp; + std::optional> meter; + std::optional eighth_position; + std::optional structure; + std::optional key; + std::optional chord; + std::vector notes; + std::vector note_end_times; + std::vector timestamp_tokens; + std::vector rhythm_tokens; + std::vector structure_tokens; + std::vector key_tokens; + std::vector chord_tokens; + std::vector melody_tokens; +}; + +int64_t sheetsage2_structure_label_count(); +int64_t sheetsage2_duration_bin_count(); +std::string sheetsage2_structure_label(int64_t index); +std::string sheetsage2_key_label(int64_t index); +std::string sheetsage2_chord_label(bool full_chord, int64_t index); +SheetSage2Note sheetsage2_note_from_pitch_duration(int pitch_id, int64_t duration_bin); + +std::string events_to_abc(const std::vector & events, double duration); +std::string events_json(const std::vector & events); + +} // namespace engine::models::sheetsage diff --git a/include/engine/models/sheetsage/runtime.h b/include/engine/models/sheetsage/runtime.h new file mode 100644 index 00000000..509a913a --- /dev/null +++ b/include/engine/models/sheetsage/runtime.h @@ -0,0 +1,70 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/models/sheetsage/types.h" + +#include +#include + +namespace engine::models::sheetsage { + +class Mert2EncoderRuntime { +public: + Mert2EncoderRuntime( + std::shared_ptr source, + core::ExecutionContext & execution, + SheetSage2DecoderConfig config = {}, + SheetSage2DecoderRuntimeOptions options = {}); + ~Mert2EncoderRuntime(); + + Mert2EncoderRuntime(const Mert2EncoderRuntime &) = delete; + Mert2EncoderRuntime & operator=(const Mert2EncoderRuntime &) = delete; + Mert2EncoderRuntime(Mert2EncoderRuntime &&) noexcept; + Mert2EncoderRuntime & operator=(Mert2EncoderRuntime &&) noexcept; + + std::vector encode_mel( + const std::vector & normalized_mel, + int64_t mel_frames); + void prepare(int64_t batch, int64_t mel_frames); + void release_runtime_graphs(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +class SheetSage2DecoderRuntime { +public: + SheetSage2DecoderRuntime( + std::shared_ptr source, + core::ExecutionContext & execution, + SheetSage2DecoderConfig config = {}, + SheetSage2DecoderRuntimeOptions options = {}); + ~SheetSage2DecoderRuntime(); + + SheetSage2DecoderRuntime(const SheetSage2DecoderRuntime &) = delete; + SheetSage2DecoderRuntime & operator=(const SheetSage2DecoderRuntime &) = delete; + SheetSage2DecoderRuntime(SheetSage2DecoderRuntime &&) noexcept; + SheetSage2DecoderRuntime & operator=(SheetSage2DecoderRuntime &&) noexcept; + + std::vector decode_logits( + const std::vector & mixed_encoder_state, + int64_t memory_steps, + const std::vector & decoder_input_ids); + void reset_cached_decode( + const std::vector & mixed_encoder_state, + int64_t memory_steps, + int64_t cache_steps); + std::vector prefill_cached_decode(const std::vector & token_ids); + std::vector decode_cached_step(int32_t token); + + void prepare(int64_t batch, int64_t memory_steps, int64_t decoder_steps); + void release_runtime_graphs(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::sheetsage diff --git a/include/engine/models/sheetsage/session.h b/include/engine/models/sheetsage/session.h new file mode 100644 index 00000000..2ce078d4 --- /dev/null +++ b/include/engine/models/sheetsage/session.h @@ -0,0 +1,56 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/audio/dsp.h" +#include "engine/models/sheetsage/runtime.h" +#include "engine/models/sheetsage/audio_frontend.h" +#include "engine/framework/model_spec/metadata.h" +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/session_base.h" + +#include +#include +#include + +namespace engine::models::sheetsage { + +struct SheetSage2Assets { + assets::ResourceBundle resources; + SheetSage2DecoderConfig config; + std::shared_ptr weights; + engine::audio::SparseMelFilterbank mel_filterbank; + std::vector mel_mean; + std::vector mel_std; + std::vector stft_window; +}; + +std::shared_ptr load_sheetsage2_assets(const std::filesystem::path & model_path); +std::shared_ptr make_sheetsage2_loader(); + +class SheetSage2Session final + : public runtime::RuntimeSessionBase + , public runtime::IOfflineVoiceTaskSession { +public: + SheetSage2Session( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract); + ~SheetSage2Session() override; + + std::string family() const override; + runtime::VoiceTaskKind task_kind() const override; + runtime::RunMode run_mode() const override; + void prepare(const runtime::SessionPreparationRequest & request) override; + runtime::TaskResult run(const runtime::TaskRequest & request) override; + +private: + runtime::TaskSpec task_; + std::shared_ptr assets_; + std::shared_ptr contract_; + Mert2EncoderRuntime encoder_; + SheetSage2AudioFrontend audio_frontend_; + SheetSage2DecoderRuntime decoder_; +}; + +} // namespace engine::models::sheetsage diff --git a/include/engine/models/sheetsage/types.h b/include/engine/models/sheetsage/types.h new file mode 100644 index 00000000..15fdfee9 --- /dev/null +++ b/include/engine/models/sheetsage/types.h @@ -0,0 +1,42 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" + +#include +#include + +namespace engine::models::sheetsage { + +struct SheetSage2DecoderConfig { + int64_t vocab_size = 31678; + int64_t hidden_size = 512; + int64_t encoder_hidden_size = 1024; + int64_t intermediate_size = 2048; + int64_t decoder_layers = 6; + int64_t num_attention_heads = 8; + int64_t max_position_embeddings = 5120; + int64_t pad_token_id = 1; + float layer_norm_eps = 1.0e-5F; + int64_t encoder_layers = 24; + int64_t encoder_attention_heads = 16; + int64_t encoder_intermediate_size = 4096; + int64_t mel_bins = 128; + int64_t sampling_rate = 24000; + int64_t n_fft = 2048; + int64_t win_length = 2048; + int64_t hop_length = 240; + int64_t input_audio_length_samples = 7200000; + int64_t convnext_kernel_size = 7; + int64_t conformer_conv_kernel_size = 31; + float encoder_layer_norm_eps = 1.0e-5F; + float subsampling_layer_norm_eps = 1.0e-6F; + float rotary_embedding_base = 10000.0F; +}; + +struct SheetSage2DecoderRuntimeOptions { + size_t graph_arena_bytes = 1536ull * 1024ull * 1024ull; + size_t weight_context_bytes = 1024ull * 1024ull * 1024ull; + assets::TensorStorageType weight_storage_type = assets::TensorStorageType::Native; +}; + +} // namespace engine::models::sheetsage diff --git a/include/engine/models/yue2/ar_runtime.h b/include/engine/models/yue2/ar_runtime.h new file mode 100644 index 00000000..20638412 --- /dev/null +++ b/include/engine/models/yue2/ar_runtime.h @@ -0,0 +1,67 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/core/module.h" +#include "engine/framework/modules/transformers/qwen_causal_decode_runtime.h" +#include "engine/framework/runtime/kv_cache.h" +#include "engine/models/yue2/assets.h" +#include "engine/models/yue2/types.h" + +#include +#include +#include +#include + +namespace engine::models::yue2 { + +struct Yue2ArSamplingWindow { + int32_t begin = 0; + int32_t end = 0; + int32_t stop_token = 0; + int64_t min_tokens = 0; + int64_t max_tokens = 0; + Yue2SamplingConfig sampling; +}; + +struct Yue2ArDevicePrefixState { + int64_t current_end = 0; + std::vector keys; + std::vector values; +}; + +class Yue2ArRuntime { +public: + Yue2ArRuntime( + core::ExecutionContext & execution, + std::shared_ptr assets, + assets::TensorStorageType weight_type, + size_t weight_context_bytes, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes); + ~Yue2ArRuntime(); + + std::vector generate( + const std::vector & prefix, + const Yue2ArSamplingWindow & window, + uint64_t seed); + + std::vector generate_cfg( + const std::vector & positive_prefix, + const std::vector & negative_prefix, + const Yue2ArSamplingWindow & window, + float guidance_scale, + uint64_t seed); + + runtime::TransformerKVState prefill_state(const std::vector & tokens); + Yue2ArDevicePrefixState prefill_device_state(const std::vector & tokens); + + void release_runtime_graphs(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::yue2 diff --git a/include/engine/models/yue2/assets.h b/include/engine/models/yue2/assets.h new file mode 100644 index 00000000..13093e2d --- /dev/null +++ b/include/engine/models/yue2/assets.h @@ -0,0 +1,21 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/models/yue2/types.h" + +#include +#include + +namespace engine::models::yue2 { + +struct Yue2Assets { + std::filesystem::path model_root; + std::filesystem::path tiktoken_path; + std::shared_ptr model_weights; + std::shared_ptr vae_weights; + Yue2Config config; +}; + +std::shared_ptr load_yue2_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::yue2 diff --git a/include/engine/models/yue2/nar_runtime.h b/include/engine/models/yue2/nar_runtime.h new file mode 100644 index 00000000..81195024 --- /dev/null +++ b/include/engine/models/yue2/nar_runtime.h @@ -0,0 +1,44 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/runtime/kv_cache.h" +#include "engine/models/yue2/ar_runtime.h" +#include "engine/models/yue2/assets.h" +#include "engine/models/yue2/types.h" + +#include +#include +#include +#include + +namespace engine::models::yue2 { + +class Yue2NarRuntime { +public: + Yue2NarRuntime( + core::ExecutionContext & execution, + std::shared_ptr assets, + assets::TensorStorageType weight_type, + size_t weight_context_bytes, + size_t graph_arena_bytes); + ~Yue2NarRuntime(); + + std::vector synthesize( + const std::vector & prefix, + const std::vector & codec, + const std::function &)> & prefill_state, + const std::vector & noise, + uint64_t seed, + int64_t ode_steps, + int64_t context); + + void release_runtime_graphs(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::yue2 diff --git a/include/engine/models/yue2/pipeline.h b/include/engine/models/yue2/pipeline.h new file mode 100644 index 00000000..ee3b2f75 --- /dev/null +++ b/include/engine/models/yue2/pipeline.h @@ -0,0 +1,46 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/codecs/oobleck_audio_vae_runtime.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/runtime/session.h" +#include "engine/models/yue2/assets.h" +#include "engine/models/yue2/request.h" +#include "engine/models/yue2/tokenizer_text.h" + +#include +#include + +namespace engine::models::yue2 { + +class Yue2PipelineRuntime { +public: + Yue2PipelineRuntime( + core::ExecutionContext & execution, + std::shared_ptr assets, + assets::TensorStorageType model_weight_type, + assets::TensorStorageType vae_weight_type, + size_t model_weight_context_bytes, + size_t vae_weight_context_bytes, + size_t ar_prefill_graph_arena_bytes, + size_t ar_decode_graph_arena_bytes, + size_t nar_graph_arena_bytes, + size_t vae_graph_arena_bytes); + ~Yue2PipelineRuntime(); + + Yue2Plan plan(const Yue2Request & request); + Yue2SemanticResult generate_semantic(const Yue2Request & request, Yue2Plan plan); + std::vector synthesize_latents( + const Yue2SemanticResult & semantic, + const Yue2GenerationConfig & generation, + uint64_t seed); + runtime::AudioBuffer decode_audio(const std::vector & latents, int64_t frames); + runtime::AudioBuffer run(const Yue2Request & request); + void release_runtime_graphs(); + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::yue2 diff --git a/include/engine/models/yue2/request.h b/include/engine/models/yue2/request.h new file mode 100644 index 00000000..746eff01 --- /dev/null +++ b/include/engine/models/yue2/request.h @@ -0,0 +1,13 @@ +#pragma once + +#include "engine/framework/runtime/session.h" +#include "engine/models/yue2/types.h" + +namespace engine::models::yue2 { + +Yue2Request parse_yue2_request(const runtime::TaskRequest & request, const Yue2GenerationConfig & defaults); +Yue2Request parse_yue2_preparation_request( + const runtime::SessionPreparationRequest & request, + const Yue2GenerationConfig & defaults); + +} // namespace engine::models::yue2 diff --git a/include/engine/models/yue2/session.h b/include/engine/models/yue2/session.h new file mode 100644 index 00000000..6fd9bb43 --- /dev/null +++ b/include/engine/models/yue2/session.h @@ -0,0 +1,36 @@ +#pragma once + +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/models/yue2/assets.h" +#include "engine/models/yue2/pipeline.h" + +#include + +namespace engine::models::yue2 { + +class Yue2Session final + : public runtime::RuntimeSessionBase + , public runtime::IOfflineVoiceTaskSession { +public: + Yue2Session( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets); + ~Yue2Session() override; + + std::string family() const override; + runtime::VoiceTaskKind task_kind() const override; + runtime::RunMode run_mode() const override; + void prepare(const runtime::SessionPreparationRequest & request) override; + runtime::TaskResult run(const runtime::TaskRequest & request) override; + +private: + runtime::TaskSpec task_; + std::shared_ptr assets_; + std::unique_ptr pipeline_; +}; + +std::shared_ptr make_yue2_loader(); + +} // namespace engine::models::yue2 diff --git a/include/engine/models/yue2/tokenizer_text.h b/include/engine/models/yue2/tokenizer_text.h new file mode 100644 index 00000000..2119fd04 --- /dev/null +++ b/include/engine/models/yue2/tokenizer_text.h @@ -0,0 +1,25 @@ +#pragma once + +#include +#include +#include +#include +#include + +namespace llama_tokenizer_vendor { +struct BpeVocabulary; +} // namespace llama_tokenizer_vendor + +namespace engine::models::yue2 { + +class Yue2TextTokenizer { +public: + explicit Yue2TextTokenizer(const std::filesystem::path & vocab_path); + + std::vector encode(const std::string & text) const; + +private: + std::shared_ptr vocab_; +}; + +} // namespace engine::models::yue2 diff --git a/include/engine/models/yue2/types.h b/include/engine/models/yue2/types.h new file mode 100644 index 00000000..a3c9f576 --- /dev/null +++ b/include/engine/models/yue2/types.h @@ -0,0 +1,105 @@ +#pragma once + +#include +#include +#include + +namespace engine::models::yue2 { + +constexpr int32_t kEodToken = 151643; +constexpr int32_t kAbcStartToken = 151847; +constexpr int32_t kAbcEndToken = 151848; +constexpr int32_t kMusicStartToken = 151851; +constexpr int32_t kMusicEndToken = 151852; +constexpr int32_t kCodecOffset = 151853; +constexpr int32_t kCodecSize = 32768; +constexpr int32_t kVocabSize = 184704; +constexpr int64_t kContextTokens = 24576; + +struct Yue2ModelConfig { + int64_t hidden_size = 2048; + int64_t layers = 28; + int64_t attention_heads = 16; + int64_t kv_heads = 8; + int64_t head_dim = 128; + int64_t intermediate_size = 6144; + int64_t vocab_size = kVocabSize; + int64_t max_position_embeddings = kContextTokens; + int64_t latent_dim = 64; + int64_t max_latent_frames = kContextTokens; + float rms_norm_eps = 1.0e-6F; + float rope_theta = 1000000.0F; + float timestep_shift = 1.0F; +}; + +struct Yue2VaeConfig { + int sample_rate = 48000; + int64_t channels = 2; + int64_t latent_dim = 64; + int64_t encoder_latent_dim = 128; + int64_t downsampling_ratio = 1920; + int64_t decode_core_frames = 1024; + int64_t decode_halo_frames = 16; +}; + +struct Yue2SamplingConfig { + float temperature = 1.0F; + float top_p = 0.95F; + int64_t top_k = 100; + float repetition_penalty = 1.2F; + int64_t penalty_window = 50; + int64_t min_tokens = 200; + int64_t max_tokens = 9000; +}; + +struct Yue2GenerationConfig { + Yue2SamplingConfig abc; + Yue2SamplingConfig semantic; + int64_t ode_steps = 32; + int64_t context = kContextTokens; +}; + +struct Yue2Config { + Yue2ModelConfig model; + Yue2VaeConfig vae; + Yue2GenerationConfig generation; +}; + +enum class Yue2CotMode { + Off, + Melody, + Full, +}; + +struct Yue2Request { + std::string style; + std::string lyrics; + Yue2CotMode cot = Yue2CotMode::Full; + std::string abc; + std::vector semantic_codes; + std::vector nar_noise; + uint64_t seed = 831001; + float cfg_scale = -1.0F; + Yue2GenerationConfig generation; +}; + +struct Yue2Plan { + Yue2CotMode cot = Yue2CotMode::Full; + std::string abc; + std::vector abc_ids; + std::vector prefix; + bool truncated = false; +}; + +struct Yue2SemanticResult { + Yue2Plan plan; + std::vector tokens; + bool truncated = false; +}; + +const char * cot_mode_name(Yue2CotMode mode) noexcept; +Yue2CotMode parse_cot_mode(const std::string & value); +const char * cot_instruction(Yue2CotMode mode) noexcept; +float request_guidance_scale(const Yue2Request & request) noexcept; + +} // namespace engine::models::yue2 diff --git a/model_specs/sheetsage2.json b/model_specs/sheetsage2.json new file mode 100644 index 00000000..7edcc031 --- /dev/null +++ b/model_specs/sheetsage2.json @@ -0,0 +1,113 @@ +{ + "schema_version": 1, + "family": "sheetsage2", + "display_name": "SheetSage2", + "description": "SheetSage2 audio-to-symbolic transcription. This native path consumes an input recording from a self-contained GGUF and emits an ABC score artifact.", + "category": "audio_tools", + "status": "experimental", + "tasks": [ + "midi" + ], + "modes": [ + "offline" + ], + "languages": [ + "music" + ], + "runtime": { + "tags": [ + "gguf", + "cuda" + ] + }, + "capabilities": { + "midi": [ + "midi_artifact" + ] + }, + "options": { + "request": [ + { + "name": "max_tokens", + "type": "int", + "description": "Maximum total decoder sequence length; default follows the embedded model context.", + "required": false, + "min": 1, + "default": 5120 + } + ], + "session": [ + { + "name": "weight_type", + "type": "enum", + "description": "Decoder weight storage type; default native.", + "preset": "weight_type_full", + "required": false, + "default": "native" + }, + { + "name": "weight_context_mb", + "type": "int", + "description": "Weight context arena size in MiB; default 1024.", + "required": false, + "min": 1, + "default": 1024 + }, + { + "name": "decoder_graph_arena_mb", + "type": "int", + "description": "Encoder/decoder graph arena size in MiB; default 1536.", + "required": false, + "min": 1, + "default": 1536 + } + ], + "load": [] + }, + "package_defaults": { + "download": { + "kind": "huggingface_snapshot", + "repo": "audio-cpp/SheetSage2-GGUF", + "revision": "main", + "gated": false + } + }, + "packages": [ + { + "id": "sheetsage2_orig", + "display_name": "SheetSage2 Original-Dtype GGUF", + "default": true, + "format": "gguf", + "precision": "orig", + "target_directory": "SheetSage2-GGUF", + "files": [ + "sheetsage2-orig.gguf" + ] + } + ], + "dependencies": [], + "ui": { + "recommended_package": "sheetsage2_orig", + "tags": [ + "Music", + "MIDI", + "GGUF" + ], + "docs": [] + }, + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "config": "model:config.json" + }, + "tensors": { + "weights": "weights:" + } + } + ] +} diff --git a/model_specs/yue2.json b/model_specs/yue2.json new file mode 100644 index 00000000..9bf293ca --- /dev/null +++ b/model_specs/yue2.json @@ -0,0 +1,444 @@ +{ + "schema_version": 1, + "family": "yue2", + "display_name": "YuE2", + "description": "YuE2 music generation model with symbolic ABC planning, semantic codec generation, NAR acoustic flow synthesis, and Oobleck VAE decode.", + "category": "audio_generation", + "status": "experimental", + "tasks": [ + "music" + ], + "modes": [ + "offline" + ], + "languages": [ + "en" + ], + "capabilities": { + "music": [ + "lyrics", + "style_control" + ] + }, + "runtime": { + "tags": [ + "gguf" + ] + }, + "dependencies": [], + "ui": { + "recommended_package": "yue2_main_q8_0", + "tags": [ + "Music", + "GGUF" + ], + "docs": [ + "docs/models/yue2.md", + "docs/music_generation.md", + "docs/gguf.md" + ] + }, + "options": { + "request": [ + { + "name": "style", + "type": "string", + "description": "Song style/tags.", + "required": true + }, + { + "name": "lyrics", + "type": "string", + "description": "Lyrics text. If omitted, the CLI text input is used.", + "required": true + }, + { + "name": "cot", + "type": "enum", + "values": [ + "off", + "melody", + "full" + ], + "description": "Symbolic planning route. off skips ABC generation; melody/full generate or consume ABC before music tokens.", + "required": false, + "default": "full" + }, + { + "name": "abc", + "type": "string", + "description": "External ABC score text for melody/full routes.", + "required": false + }, + { + "name": "abc_file", + "type": "string", + "description": "Path to an external ABC score file for melody/full routes.", + "required": false + }, + { + "name": "semantic_codes_file", + "type": "string", + "description": "Teacher-forced raw int32 semantic codec IDs for downstream parity validation.", + "required": false + }, + { + "name": "nar_noise_file", + "type": "string", + "description": "Teacher-forced raw float32 NAR noise [frames,64] for downstream parity validation.", + "required": false + }, + { + "name": "cfg_scale", + "type": "float", + "description": "Semantic classifier-free guidance scale. Default follows the upstream route defaults.", + "required": false, + "min": 0.0, + "max": 20.0 + }, + { + "name": "seed", + "type": "int", + "description": "Generation seed.", + "required": false, + "min": 0, + "default": 831001 + }, + { + "name": "num_inference_steps", + "type": "int", + "description": "NAR midpoint ODE steps.", + "required": false, + "min": 1, + "default": 32 + }, + { + "name": "abc_temperature", + "type": "float", + "description": "ABC planner sampling temperature.", + "required": false, + "min": 0.0, + "max": 5.0 + }, + { + "name": "abc_top_p", + "type": "float", + "description": "ABC planner nucleus sampling probability.", + "required": false, + "min": 0.0, + "max": 1.0 + }, + { + "name": "abc_top_k", + "type": "int", + "description": "ABC planner top-k sampling limit.", + "required": false, + "min": 1 + }, + { + "name": "abc_repetition_penalty", + "type": "float", + "description": "ABC planner repetition penalty.", + "required": false, + "min": 0.001 + }, + { + "name": "abc_penalty_window", + "type": "int", + "description": "ABC planner repetition penalty window.", + "required": false, + "min": 1 + }, + { + "name": "abc_min_tokens", + "type": "int", + "description": "Minimum ABC planner tokens before EOS is accepted.", + "required": false, + "min": 0 + }, + { + "name": "abc_max_tokens", + "type": "int", + "description": "Maximum ABC planner tokens.", + "required": false, + "min": 1 + }, + { + "name": "semantic_temperature", + "type": "float", + "description": "Semantic codec sampling temperature.", + "required": false, + "min": 0.0, + "max": 5.0 + }, + { + "name": "semantic_top_p", + "type": "float", + "description": "Semantic codec nucleus sampling probability.", + "required": false, + "min": 0.0, + "max": 1.0 + }, + { + "name": "semantic_top_k", + "type": "int", + "description": "Semantic codec top-k sampling limit.", + "required": false, + "min": 1 + }, + { + "name": "semantic_repetition_penalty", + "type": "float", + "description": "Semantic codec repetition penalty.", + "required": false, + "min": 0.001 + }, + { + "name": "semantic_penalty_window", + "type": "int", + "description": "Semantic codec repetition penalty window.", + "required": false, + "min": 1 + }, + { + "name": "semantic_min_tokens", + "type": "int", + "description": "Minimum semantic tokens before EOS is accepted.", + "required": false, + "min": 0 + }, + { + "name": "semantic_max_tokens", + "type": "int", + "description": "Maximum semantic codec tokens.", + "required": false, + "min": 1 + } + ], + "session": [ + { + "name": "weight_type", + "type": "enum", + "values": [ + "native", + "f32", + "f16", + "bf16", + "q8_0", + "q4_0", + "q4_k" + ], + "description": "Shared weight storage type.", + "required": false, + "default": "native" + }, + { + "name": "model_weight_type", + "type": "enum", + "values": [ + "native", + "f32", + "f16", + "bf16", + "q8_0", + "q4_0", + "q4_k" + ], + "description": "YuE2 MoT weight storage type.", + "required": false, + "default": "native" + }, + { + "name": "model_gguf", + "type": "string", + "description": "Yue2 main AR/NAR component GGUF file relative to the model root.", + "required": false, + "default": "yue2-3b-q8_0.gguf" + }, + { + "name": "vae_gguf", + "type": "string", + "description": "Yue2 VAE component GGUF file relative to the model root.", + "required": false, + "default": "yue2-vae-f16.gguf" + }, + { + "name": "vae_weight_type", + "type": "enum", + "values": [ + "native", + "f32", + "f16", + "bf16", + "q8_0", + "q4_0", + "q4_k" + ], + "description": "YuE2 VAE weight storage type.", + "required": false, + "default": "native" + }, + { + "name": "model_weight_context_mb", + "type": "int", + "description": "YuE2 MoT weight context size in MiB.", + "required": false, + "min": 1, + "default": 6144 + }, + { + "name": "vae_weight_context_mb", + "type": "int", + "description": "YuE2 VAE weight context size in MiB.", + "required": false, + "min": 1, + "default": 1536 + }, + { + "name": "ar_prefill_graph_arena_mb", + "type": "int", + "description": "AR prefill graph arena size in MiB.", + "required": false, + "min": 1, + "default": 4096 + }, + { + "name": "ar_decode_graph_arena_mb", + "type": "int", + "description": "AR one-token decode graph arena size in MiB.", + "required": false, + "min": 1, + "default": 1536 + }, + { + "name": "nar_graph_arena_mb", + "type": "int", + "description": "NAR acoustic flow graph arena size in MiB.", + "required": false, + "min": 1, + "default": 6144 + }, + { + "name": "vae_graph_arena_mb", + "type": "int", + "description": "VAE decode graph arena size in MiB.", + "required": false, + "min": 1, + "default": 1536 + } + ], + "load": [] + }, + "package_defaults": { + "download": { + "kind": "huggingface_snapshot", + "repo": "audio-cpp/Yue2-3B-GGUF", + "revision": "main", + "gated": false + } + }, + "packages": [ + { + "id": "yue2_main_q8_0", + "display_name": "Yue2 3B Main Q8_0", + "default": true, + "format": "gguf", + "precision": "q8_0", + "target_directory": "Yue2-3B-GGUF", + "files": [ + "sidecars/yue2-model-config.json", + "sidecars/yue2-generation-config.json", + "sidecars/yue2-qwen.tiktoken", + "sidecars/yue2-vae-config.json", + "yue2-3b-q8_0.gguf" + ] + }, + { + "id": "yue2_main_bf16", + "display_name": "Yue2 3B Main BF16", + "format": "gguf", + "precision": "bf16", + "target_directory": "Yue2-3B-GGUF", + "files": [ + "sidecars/yue2-model-config.json", + "sidecars/yue2-generation-config.json", + "sidecars/yue2-qwen.tiktoken", + "sidecars/yue2-vae-config.json", + "yue2-3b-bf16.gguf" + ] + }, + { + "id": "yue2_main_q4_0", + "display_name": "Yue2 3B Main Q4_0", + "format": "gguf", + "precision": "q4_0", + "target_directory": "Yue2-3B-GGUF", + "files": [ + "sidecars/yue2-model-config.json", + "sidecars/yue2-generation-config.json", + "sidecars/yue2-qwen.tiktoken", + "sidecars/yue2-vae-config.json", + "yue2-3b-q4_0.gguf" + ] + }, + { + "id": "yue2_vae_f16", + "display_name": "Yue2 VAE F16", + "format": "gguf", + "precision": "f16", + "target_directory": "Yue2-3B-GGUF", + "files": [ + "sidecars/yue2-model-config.json", + "sidecars/yue2-generation-config.json", + "sidecars/yue2-qwen.tiktoken", + "sidecars/yue2-vae-config.json", + "yue2-vae-f16.gguf" + ] + }, + { + "id": "yue2_vae_f32", + "display_name": "Yue2 VAE F32", + "format": "gguf", + "precision": "f32", + "target_directory": "Yue2-3B-GGUF", + "files": [ + "sidecars/yue2-model-config.json", + "sidecars/yue2-generation-config.json", + "sidecars/yue2-qwen.tiktoken", + "sidecars/yue2-vae-config.json", + "yue2-vae-f32.gguf" + ] + } + ], + "sources": [ + { + "format": "gguf", + "roots": { + "model": "." + }, + "files": { + "model_config": "model:sidecars/yue2-model-config.json", + "generation_config": "model:sidecars/yue2-generation-config.json", + "tiktoken": "model:sidecars/yue2-qwen.tiktoken", + "vae_config": "model:sidecars/yue2-vae-config.json" + } + }, + { + "format": "safetensors", + "roots": { + "model": "YuE2-3B", + "vae": "YuE2-Vae" + }, + "files": { + "model_config": "model:config.json", + "generation_config": "model:yue2_generation_config.json", + "tiktoken": "model:qwen.tiktoken", + "vae_config": "vae:config.json" + }, + "tensors": { + "model_weights": "model:model.safetensors", + "vae_weights": "vae:model.safetensors" + } + } + ] +} diff --git a/src/framework/audio/conversion.cpp b/src/framework/audio/conversion.cpp index 0ee640fc..a002639b 100644 --- a/src/framework/audio/conversion.cpp +++ b/src/framework/audio/conversion.cpp @@ -143,6 +143,23 @@ std::vector convert_interleaved_audio_to_mono_linear_resampled( target_sample_rate_hz); } +std::vector convert_interleaved_audio_to_mono_torchaudio_sinc_hann_resampled( + const std::vector & interleaved_samples, + int sample_rate_hz, + int channel_count, + int target_sample_rate_hz, + const TorchaudioSincHannResampleOptions & options, + MonoMixAccumulation accumulation) { + if (sample_rate_hz <= 0 || target_sample_rate_hz <= 0) { + throw std::runtime_error("audio sample rates must be positive"); + } + auto mono = mixdown_interleaved_to_mono_average(interleaved_samples, channel_count, accumulation); + if (sample_rate_hz != target_sample_rate_hz) { + mono = resample_mono_torchaudio_sinc_hann(mono, sample_rate_hz, target_sample_rate_hz, options); + } + return mono; +} + std::vector read_wav_f32_as_mono_linear_resampled( const std::filesystem::path & path, int target_sample_rate_hz) { diff --git a/src/framework/codecs/oobleck_audio_vae_runtime.cpp b/src/framework/codecs/oobleck_audio_vae_runtime.cpp new file mode 100644 index 00000000..c726f633 --- /dev/null +++ b/src/framework/codecs/oobleck_audio_vae_runtime.cpp @@ -0,0 +1,752 @@ +#include "engine/framework/codecs/oobleck_audio_vae_runtime.h" + +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::codecs { +namespace { + +using Clock = std::chrono::steady_clock; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct OobleckResidualUnitWeights { + modules::SnakeBeta1dWeights snake_0; + modules::Conv1dWeights conv_1; + modules::SnakeBeta1dWeights snake_2; + modules::Conv1dWeights conv_3; +}; + +struct OobleckDecoderBlockWeights { + modules::SnakeBeta1dWeights snake_0; + modules::ConvTranspose1dWeights upsample; + OobleckResidualUnitWeights residual_2; + OobleckResidualUnitWeights residual_3; + OobleckResidualUnitWeights residual_4; +}; + +struct OobleckEncoderBlockWeights { + OobleckResidualUnitWeights residual_0; + OobleckResidualUnitWeights residual_1; + OobleckResidualUnitWeights residual_2; + modules::SnakeBeta1dWeights snake_3; + modules::Conv1dWeights downsample; +}; + +struct OobleckAudioVaeWeights { + std::shared_ptr store; + modules::Conv1dWeights encoder_in_conv; + std::vector encoder_blocks; + modules::SnakeBeta1dWeights encoder_final_snake; + modules::Conv1dWeights encoder_out_conv; + modules::Conv1dWeights decoder_in_conv; + std::vector decoder_blocks; + modules::SnakeBeta1dWeights decoder_final_snake; + modules::Conv1dWeights decoder_out_conv; +}; + +std::vector channel_plan(const OobleckAudioVaeConfig & config) { + if (config.c_mults.empty()) { + throw std::runtime_error("Oobleck audio VAE c_mults must not be empty"); + } + std::vector channels; + channels.reserve(config.c_mults.size() + 1); + channels.push_back(1); + channels.insert(channels.end(), config.c_mults.begin(), config.c_mults.end()); + for (int64_t & value : channels) { + value *= config.channels; + } + return channels; +} + +void validate_config(const OobleckAudioVaeConfig & config) { + if (config.sample_rate <= 0 || config.audio_channels <= 0 || config.channels <= 0 || + config.encoder_latent_dim <= 0 || config.decoder_latent_dim <= 0) { + throw std::runtime_error("Oobleck audio VAE config dimensions must be positive"); + } + if (config.c_mults.empty() || config.c_mults.size() != config.strides.size()) { + throw std::runtime_error("Oobleck audio VAE c_mults and strides must have matching positive sizes"); + } + if (!config.use_snake) { + throw std::runtime_error("Oobleck audio VAE framework runtime currently supports SnakeBeta checkpoints"); + } + for (const int64_t stride : config.strides) { + if (stride <= 0) { + throw std::runtime_error("Oobleck audio VAE strides must be positive"); + } + } +} + +std::string join_name(const std::string & prefix, const std::string & name) { + return prefix.empty() ? name : prefix + "." + name; +} + +std::vector effective_weight_norm_conv( + const assets::TensorSource & source, + const std::string & prefix, + const std::vector & shape) { + if (shape.size() != 3) { + throw std::runtime_error("Oobleck audio VAE weight-norm conv shape must be rank 3"); + } + auto v = source.require_f32(prefix + ".weight_v", shape); + auto g = source.require_f32(prefix + ".weight_g", {shape[0], 1, 1}); + const int64_t outer = shape[0]; + const int64_t inner = shape[1] * shape[2]; + for (int64_t o = 0; o < outer; ++o) { + double norm_sq = 0.0; + for (int64_t i = 0; i < inner; ++i) { + const float value = v[static_cast(o * inner + i)]; + norm_sq += static_cast(value) * static_cast(value); + } + const float scale = g[static_cast(o)] / std::sqrt(static_cast(std::max(norm_sq, 1.0e-24))); + for (int64_t i = 0; i < inner; ++i) { + v[static_cast(o * inner + i)] *= scale; + } + } + return v; +} + +modules::Conv1dWeights load_wn_conv1d( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t out_channels, + int64_t in_channels, + int64_t kernel, + bool use_bias) { + modules::Conv1dWeights weights; + if (source.has_tensor(prefix + ".weight")) { + weights.weight = store.load_tensor( + source, + prefix + ".weight", + storage_type, + {out_channels, in_channels, kernel}); + } else { + weights.weight = store.make_from_f32( + core::TensorShape::from_dims({out_channels, in_channels, kernel}), + storage_type, + effective_weight_norm_conv(source, prefix, {out_channels, in_channels, kernel})); + } + if (use_bias) { + weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_channels}); + } + return weights; +} + +modules::ConvTranspose1dWeights load_wn_conv_transpose1d( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t in_channels, + int64_t out_channels, + int64_t kernel, + bool use_bias) { + modules::ConvTranspose1dWeights weights; + if (source.has_tensor(prefix + ".weight")) { + weights.weight = store.load_tensor( + source, + prefix + ".weight", + storage_type, + {in_channels, out_channels, kernel}); + } else { + weights.weight = store.make_from_f32( + core::TensorShape::from_dims({in_channels, out_channels, kernel}), + storage_type, + effective_weight_norm_conv(source, prefix, {in_channels, out_channels, kernel})); + } + if (use_bias) { + weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_channels}); + } + return weights; +} + +modules::SnakeBeta1dWeights load_snake( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t channels) { + return { + store.load_f32_tensor(source, prefix + ".alpha", {channels}), + store.load_f32_tensor(source, prefix + ".beta", {channels}), + }; +} + +OobleckResidualUnitWeights load_residual_unit( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t channels) { + return { + load_snake(store, source, prefix + ".layers.0", channels), + load_wn_conv1d(store, source, prefix + ".layers.1", storage_type, channels, channels, 7, true), + load_snake(store, source, prefix + ".layers.2", channels), + load_wn_conv1d(store, source, prefix + ".layers.3", storage_type, channels, channels, 1, true), + }; +} + +OobleckEncoderBlockWeights load_encoder_block( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t in_channels, + int64_t out_channels, + int64_t stride) { + return { + load_residual_unit(store, source, prefix + ".layers.0", storage_type, in_channels), + load_residual_unit(store, source, prefix + ".layers.1", storage_type, in_channels), + load_residual_unit(store, source, prefix + ".layers.2", storage_type, in_channels), + load_snake(store, source, prefix + ".layers.3", in_channels), + load_wn_conv1d(store, source, prefix + ".layers.4", storage_type, out_channels, in_channels, 2 * stride, true), + }; +} + +OobleckDecoderBlockWeights load_decoder_block( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t in_channels, + int64_t out_channels, + int64_t stride) { + return { + load_snake(store, source, prefix + ".layers.0", in_channels), + load_wn_conv_transpose1d(store, source, prefix + ".layers.1", storage_type, in_channels, out_channels, 2 * stride, true), + load_residual_unit(store, source, prefix + ".layers.2", storage_type, out_channels), + load_residual_unit(store, source, prefix + ".layers.3", storage_type, out_channels), + load_residual_unit(store, source, prefix + ".layers.4", storage_type, out_channels), + }; +} + +OobleckAudioVaeWeights load_weights( + const assets::TensorSource & source, + const OobleckAudioVaeConfig & config, + ggml_backend_t backend, + core::BackendType backend_type, + const OobleckAudioVaeRuntimeOptions & options) { + const auto total_start = Clock::now(); + const auto channels = channel_plan(config); + OobleckAudioVaeWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "framework.oobleck_audio_vae.weights", + options.weight_context_bytes); + const auto bind_start = Clock::now(); + weights.encoder_in_conv = load_wn_conv1d( + *weights.store, + source, + join_name(config.encoder_prefix, "layers.0"), + options.weight_storage_type, + channels.front(), + config.audio_channels, + 7, + true); + weights.encoder_blocks.reserve(config.strides.size()); + for (size_t i = 0; i < config.strides.size(); ++i) { + weights.encoder_blocks.push_back(load_encoder_block( + *weights.store, + source, + join_name(config.encoder_prefix, "layers." + std::to_string(i + 1)), + options.weight_storage_type, + channels[i], + channels[i + 1], + config.strides[i])); + } + weights.encoder_final_snake = load_snake( + *weights.store, + source, + join_name(config.encoder_prefix, "layers." + std::to_string(config.strides.size() + 1)), + channels.back()); + weights.encoder_out_conv = load_wn_conv1d( + *weights.store, + source, + join_name(config.encoder_prefix, "layers." + std::to_string(config.strides.size() + 2)), + options.weight_storage_type, + config.encoder_latent_dim, + channels.back(), + 3, + true); + weights.decoder_in_conv = load_wn_conv1d( + *weights.store, + source, + join_name(config.decoder_prefix, "layers.0"), + options.weight_storage_type, + channels.back(), + config.decoder_latent_dim, + 7, + true); + weights.decoder_blocks.reserve(config.strides.size()); + for (int64_t block = static_cast(config.strides.size()); block > 0; --block) { + const int64_t layer_index = static_cast(config.strides.size()) - block + 1; + weights.decoder_blocks.push_back(load_decoder_block( + *weights.store, + source, + join_name(config.decoder_prefix, "layers." + std::to_string(layer_index)), + options.weight_storage_type, + channels[static_cast(block)], + channels[static_cast(block - 1)], + config.strides[static_cast(block - 1)])); + } + weights.decoder_final_snake = load_snake( + *weights.store, + source, + join_name(config.decoder_prefix, "layers." + std::to_string(config.strides.size() + 1)), + channels.front()); + weights.decoder_out_conv = load_wn_conv1d( + *weights.store, + source, + join_name(config.decoder_prefix, "layers." + std::to_string(config.strides.size() + 2)), + options.weight_storage_type, + config.audio_channels, + channels.front(), + 7, + false); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.weights_bind_ms", engine::debug::elapsed_ms(bind_start)); + const auto upload_start = Clock::now(); + weights.store->upload(); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.weights_upload_ms", engine::debug::elapsed_ms(upload_start)); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.weights_total_ms", engine::debug::elapsed_ms(total_start)); + return weights; +} + +core::TensorValue snake_beta( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::SnakeBeta1dWeights & weights, + int64_t channels, + bool logscale) { + return modules::SnakeBeta1dModule({channels, logscale}).build(ctx, input, weights); +} + +core::TensorValue residual_unit( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const OobleckResidualUnitWeights & weights, + int64_t channels, + int64_t dilation, + bool logscale) { + auto hidden = snake_beta(ctx, input, weights.snake_0, channels, logscale); + hidden = modules::Conv1dModule({channels, channels, 7, 1, static_cast(dilation * 3), static_cast(dilation), true}) + .build(ctx, hidden, weights.conv_1); + hidden = snake_beta(ctx, hidden, weights.snake_2, channels, logscale); + hidden = modules::Conv1dModule({channels, channels, 1, 1, 0, 1, true}).build(ctx, hidden, weights.conv_3); + return modules::AddModule{}.build(ctx, input, hidden); +} + +core::TensorValue encoder_block( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const OobleckEncoderBlockWeights & weights, + int64_t in_channels, + int64_t out_channels, + int64_t stride, + bool logscale) { + auto hidden = residual_unit(ctx, input, weights.residual_0, in_channels, 1, logscale); + hidden = residual_unit(ctx, hidden, weights.residual_1, in_channels, 3, logscale); + hidden = residual_unit(ctx, hidden, weights.residual_2, in_channels, 9, logscale); + hidden = snake_beta(ctx, hidden, weights.snake_3, in_channels, logscale); + return modules::Conv1dModule({ + in_channels, + out_channels, + 2 * stride, + static_cast(stride), + static_cast(std::ceil(static_cast(stride) / 2.0F)), + 1, + true}) + .build(ctx, hidden, weights.downsample); +} + +core::TensorValue decoder_block( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const OobleckDecoderBlockWeights & weights, + int64_t in_channels, + int64_t out_channels, + int64_t stride, + bool logscale) { + auto hidden = snake_beta(ctx, input, weights.snake_0, in_channels, logscale); + hidden = modules::ConvTranspose1dModule({ + in_channels, + out_channels, + 2 * stride, + static_cast(stride), + 0, + 1, + true}) + .build(ctx, hidden, weights.upsample); + hidden = modules::SliceModule({ + 2, + static_cast(std::ceil(static_cast(stride) / 2.0F)), + input.shape.dims[2] * stride + stride - 2 * static_cast(std::ceil(static_cast(stride) / 2.0F))}) + .build(ctx, hidden); + hidden = residual_unit(ctx, hidden, weights.residual_2, out_channels, 1, logscale); + hidden = residual_unit(ctx, hidden, weights.residual_3, out_channels, 3, logscale); + hidden = residual_unit(ctx, hidden, weights.residual_4, out_channels, 9, logscale); + return hidden; +} + +} // namespace + +struct OobleckAudioVaeRuntime::Impl { + class EncodeGraph; + class DecodeGraph; + + Impl( + std::shared_ptr source, + core::ExecutionContext & execution, + OobleckAudioVaeConfig config, + OobleckAudioVaeRuntimeOptions options) + : source(std::move(source)), + execution(&execution), + config(std::move(config)), + options(options) { + if (!this->source) { + throw std::runtime_error("Oobleck audio VAE runtime requires tensor source"); + } + validate_config(this->config); + } + + const OobleckAudioVaeWeights & require_weights() { + if (!weights) { + const auto start = Clock::now(); + weights = std::make_unique(load_weights( + *source, + config, + execution->backend(), + execution->backend_type(), + options)); + source->release_storage(); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.require_weights_ms", engine::debug::elapsed_ms(start)); + } + return *weights; + } + + std::shared_ptr source; + core::ExecutionContext * execution = nullptr; + OobleckAudioVaeConfig config; + OobleckAudioVaeRuntimeOptions options; + std::unique_ptr weights; + std::unique_ptr encode_graph; + std::unique_ptr decode_graph; +}; + +class OobleckAudioVaeRuntime::Impl::EncodeGraph { +public: + EncodeGraph(core::ExecutionContext & execution, const OobleckAudioVaeConfig & config, const OobleckAudioVaeRuntimeOptions & options, const OobleckAudioVaeWeights & weights, int64_t frames) + : backend_(execution.backend()), + backend_type_(execution.backend_type()), + threads_(std::max(1, execution.config().threads)), + config_(config), + options_(options), + weights_(weights), + frames_(frames) { + if (backend_ == nullptr || frames_ <= 0) { + throw std::runtime_error("Oobleck audio VAE encode graph initialization failed"); + } + build(); + } + + ~EncodeGraph() { + if (backend_ != nullptr && graph_ != nullptr) { + core::release_backend_graph_resources(backend_type_, backend_, graph_); + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + } + + bool matches(int64_t frames) const noexcept { + return frames == frames_; + } + + std::vector run(const std::vector & planar_audio) const { + if (static_cast(planar_audio.size()) != config_.audio_channels * frames_) { + throw std::runtime_error("Oobleck audio VAE encode input shape mismatch"); + } + const auto total_start = Clock::now(); + const auto upload_start = Clock::now(); + core::write_tensor_f32(input_, planar_audio); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.encode.input_upload_ms", engine::debug::elapsed_ms(upload_start)); + core::set_backend_threads(backend_, threads_); + const auto compute_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(backend_, graph_, nullptr, "framework.oobleck_audio_vae.encode"); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.encode.graph_compute_ms", engine::debug::elapsed_ms(compute_start)); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Oobleck audio VAE encode graph compute failed"); + } + const auto read_start = Clock::now(); + auto out = core::read_tensor_f32(output_); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.encode.output_read_ms", engine::debug::elapsed_ms(read_start)); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.encode.total_ms", engine::debug::elapsed_ms(total_start)); + return out; + } + +private: + void build() { + const auto start = Clock::now(); + ggml_init_params params{options_.graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("Oobleck audio VAE encode ggml context initialization failed"); + } + core::ModuleBuildContext input_ctx{ctx_.get(), "framework.oobleck_audio_vae.encode.inputs", backend_type_}; + input_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, config_.audio_channels, frames_})); + ggml_set_input(input_.tensor); + core::ModuleBuildContext build_ctx{ctx_.get(), "framework.oobleck_audio_vae.encode", backend_type_}; + auto output = build_graph_output(build_ctx); + output_ = output.tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), 524288, false); + ggml_build_forward_expand(graph_, output_); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (gallocr_ == nullptr || !ggml_gallocr_reserve(gallocr_, graph_) || !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + throw std::runtime_error("Oobleck audio VAE encode backend buffer allocation failed"); + } + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.encode.graph_frames", frames_); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.encode.graph_build_ms", engine::debug::elapsed_ms(start)); + } + + core::TensorValue build_graph_output(core::ModuleBuildContext & ctx) const { + const auto channels = channel_plan(config_); + auto hidden = modules::Conv1dModule({config_.audio_channels, channels.front(), 7, 1, 3, 1, true}) + .build(ctx, input_, weights_.encoder_in_conv); + for (size_t i = 0; i < weights_.encoder_blocks.size(); ++i) { + hidden = encoder_block(ctx, hidden, weights_.encoder_blocks[i], channels[i], channels[i + 1], config_.strides[i], config_.snake_logscale); + } + hidden = snake_beta(ctx, hidden, weights_.encoder_final_snake, channels.back(), config_.snake_logscale); + return modules::Conv1dModule({channels.back(), config_.encoder_latent_dim, 3, 1, 1, 1, true}) + .build(ctx, hidden, weights_.encoder_out_conv); + } + + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + OobleckAudioVaeConfig config_; + OobleckAudioVaeRuntimeOptions options_; + const OobleckAudioVaeWeights & weights_; + int64_t frames_ = 0; + std::unique_ptr ctx_; + core::TensorValue input_; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; +}; + +class OobleckAudioVaeRuntime::Impl::DecodeGraph { +public: + DecodeGraph(core::ExecutionContext & execution, const OobleckAudioVaeConfig & config, const OobleckAudioVaeRuntimeOptions & options, const OobleckAudioVaeWeights & weights, int64_t batch, int64_t latent_frames) + : backend_(execution.backend()), + backend_type_(execution.backend_type()), + threads_(std::max(1, execution.config().threads)), + config_(config), + options_(options), + weights_(weights), + batch_(batch), + latent_frames_(latent_frames) { + if (backend_ == nullptr || batch_ <= 0 || latent_frames_ <= 0) { + throw std::runtime_error("Oobleck audio VAE decode graph initialization failed"); + } + build(); + } + + ~DecodeGraph() { + if (backend_ != nullptr && graph_ != nullptr) { + core::release_backend_graph_resources(backend_type_, backend_, graph_); + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + } + + bool matches(int64_t batch, int64_t latent_frames) const noexcept { + return batch == batch_ && latent_frames == latent_frames_; + } + + std::vector run(const std::vector & latents) const { + if (static_cast(latents.size()) != batch_ * config_.decoder_latent_dim * latent_frames_) { + throw std::runtime_error("Oobleck audio VAE latent shape mismatch"); + } + const auto total_start = Clock::now(); + const auto upload_start = Clock::now(); + core::write_tensor_f32(input_, latents); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.decode.input_upload_ms", engine::debug::elapsed_ms(upload_start)); + core::set_backend_threads(backend_, threads_); + const auto compute_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(backend_, graph_, nullptr, "framework.oobleck_audio_vae.decode"); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.decode.graph_compute_ms", engine::debug::elapsed_ms(compute_start)); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Oobleck audio VAE decode graph compute failed"); + } + const auto read_start = Clock::now(); + auto out = core::read_tensor_f32(output_); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.decode.output_read_ms", engine::debug::elapsed_ms(read_start)); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.decode.total_ms", engine::debug::elapsed_ms(total_start)); + return out; + } + +private: + void build() { + const auto start = Clock::now(); + ggml_init_params params{options_.graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("Oobleck audio VAE decode ggml context initialization failed"); + } + core::ModuleBuildContext input_ctx{ctx_.get(), "framework.oobleck_audio_vae.decode.inputs", backend_type_}; + input_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_, config_.decoder_latent_dim, latent_frames_})); + ggml_set_input(input_.tensor); + core::ModuleBuildContext build_ctx{ctx_.get(), "framework.oobleck_audio_vae.decode", backend_type_}; + auto output = build_graph_output(build_ctx); + output_ = output.tensor; + output_frames_ = output.shape.dims[2]; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), 524288, false); + ggml_build_forward_expand(graph_, output_); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (gallocr_ == nullptr || !ggml_gallocr_reserve(gallocr_, graph_) || !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + throw std::runtime_error("Oobleck audio VAE decode backend buffer allocation failed"); + } + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.decode.graph_batch", batch_); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.decode.graph_latent_frames", latent_frames_); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.decode.graph_output_frames", output_frames_); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.decode.graph_build_ms", engine::debug::elapsed_ms(start)); + } + + core::TensorValue build_graph_output(core::ModuleBuildContext & ctx) const { + const auto channels = channel_plan(config_); + auto hidden = modules::Conv1dModule({config_.decoder_latent_dim, channels.back(), 7, 1, 3, 1, true}) + .build(ctx, input_, weights_.decoder_in_conv); + for (size_t i = 0; i < weights_.decoder_blocks.size(); ++i) { + const int64_t block = static_cast(config_.strides.size() - i); + hidden = decoder_block( + ctx, + hidden, + weights_.decoder_blocks[i], + channels[static_cast(block)], + channels[static_cast(block - 1)], + config_.strides[static_cast(block - 1)], + config_.snake_logscale); + } + hidden = snake_beta(ctx, hidden, weights_.decoder_final_snake, channels.front(), config_.snake_logscale); + auto output = modules::Conv1dModule({channels.front(), config_.audio_channels, 7, 1, 3, 1, false}) + .build(ctx, hidden, weights_.decoder_out_conv); + if (config_.final_tanh) { + output = modules::TanhModule{}.build(ctx, output); + } + return output; + } + + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + OobleckAudioVaeConfig config_; + OobleckAudioVaeRuntimeOptions options_; + const OobleckAudioVaeWeights & weights_; + int64_t batch_ = 0; + int64_t latent_frames_ = 0; + int64_t output_frames_ = 0; + std::unique_ptr ctx_; + core::TensorValue input_; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; +}; + +OobleckAudioVaeRuntime::OobleckAudioVaeRuntime( + std::shared_ptr source, + core::ExecutionContext & execution, + OobleckAudioVaeConfig config, + OobleckAudioVaeRuntimeOptions options) + : impl_(std::make_unique(std::move(source), execution, std::move(config), options)) {} + +OobleckAudioVaeRuntime::~OobleckAudioVaeRuntime() = default; +OobleckAudioVaeRuntime::OobleckAudioVaeRuntime(OobleckAudioVaeRuntime &&) noexcept = default; +OobleckAudioVaeRuntime & OobleckAudioVaeRuntime::operator=(OobleckAudioVaeRuntime &&) noexcept = default; + +void OobleckAudioVaeRuntime::prepare_encode(int64_t frames) { + const auto & weights = impl_->require_weights(); + if (!impl_->encode_graph || !impl_->encode_graph->matches(frames)) { + impl_->encode_graph = std::make_unique(*impl_->execution, impl_->config, impl_->options, weights, frames); + } +} + +void OobleckAudioVaeRuntime::prepare_decode(int64_t batch, int64_t latent_frames) { + const auto & weights = impl_->require_weights(); + if (!impl_->decode_graph || !impl_->decode_graph->matches(batch, latent_frames)) { + impl_->decode_graph = std::make_unique(*impl_->execution, impl_->config, impl_->options, weights, batch, latent_frames); + } +} + +std::vector OobleckAudioVaeRuntime::encode_planar(const std::vector & planar_audio, int64_t frames) { + const auto start = Clock::now(); + prepare_encode(frames); + auto out = impl_->encode_graph->run(planar_audio); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.encode_planar_ms", engine::debug::elapsed_ms(start)); + return out; +} + +std::vector OobleckAudioVaeRuntime::decode_planar(const std::vector & latents, int64_t batch, int64_t latent_frames) { + const auto start = Clock::now(); + prepare_decode(batch, latent_frames); + auto out = impl_->decode_graph->run(latents); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.decode_planar_ms", engine::debug::elapsed_ms(start)); + return out; +} + +std::vector OobleckAudioVaeRuntime::decode(const std::vector & latents, int64_t batch, int64_t latent_frames) { + const auto total_start = Clock::now(); + auto planar = decode_planar(latents, batch, latent_frames); + const int64_t frames = static_cast(planar.size()) / (batch * impl_->config.audio_channels); + std::vector out; + out.reserve(static_cast(batch)); + const auto interleave_start = Clock::now(); + for (int64_t b = 0; b < batch; ++b) { + runtime::AudioBuffer audio; + audio.sample_rate = static_cast(impl_->config.sample_rate); + audio.channels = static_cast(impl_->config.audio_channels); + audio.samples.assign(static_cast(frames * impl_->config.audio_channels), 0.0F); + for (int64_t c = 0; c < impl_->config.audio_channels; ++c) { + for (int64_t t = 0; t < frames; ++t) { + audio.samples[static_cast(t * impl_->config.audio_channels + c)] = + planar[static_cast((b * impl_->config.audio_channels + c) * frames + t)]; + } + } + out.push_back(std::move(audio)); + } + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.decode.interleave_ms", engine::debug::elapsed_ms(interleave_start)); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.decode.audio_frames", frames); + engine::debug::timing_log_scalar("framework.oobleck_audio_vae.decode.full_ms", engine::debug::elapsed_ms(total_start)); + return out; +} + +void OobleckAudioVaeRuntime::release_runtime_graphs() { + impl_->encode_graph.reset(); + impl_->decode_graph.reset(); +} + +} // namespace engine::codecs diff --git a/src/framework/modules/activation_modules.cpp b/src/framework/modules/activation_modules.cpp index 8002d1e1..ca9725c5 100644 --- a/src/framework/modules/activation_modules.cpp +++ b/src/framework/modules/activation_modules.cpp @@ -149,6 +149,12 @@ const core::ModulePortSpec kSnakeInputs[] = { {"alpha", core::PortKind::Parameter, false}, }; +const core::ModulePortSpec kSnakeBetaInputs[] = { + {"input", core::PortKind::Activation, false}, + {"alpha", core::PortKind::Parameter, false}, + {"beta", core::PortKind::Parameter, false}, +}; + const core::ModulePortSpec kAliasFreeActivationInputs[] = { {"input", core::PortKind::Activation, false}, {"alpha", core::PortKind::Parameter, false}, @@ -168,6 +174,16 @@ const core::ModuleSchema kSnake1dSchema = { "Applies Snake activation over channel-time tensors using per-channel alpha.", }; +const core::ModuleSchema kSnakeBeta1dSchema = { + "SnakeBeta1d", + "nn.activation", + kSnakeBetaInputs, + 3, + kActivationOutputs, + 1, + "Applies SnakeBeta activation over channel-time tensors using per-channel alpha and beta.", +}; + const core::ModuleSchema kAliasFreeActivationSchema = { "AliasFreeActivation", "nn.activation", @@ -604,6 +620,55 @@ const core::ModuleSchema & Snake1dModule::static_schema() noexcept { return kSnake1dSchema; } +SnakeBeta1dModule::SnakeBeta1dModule(SnakeBeta1dConfig config) : config_(config) { + if (config_.hidden_size <= 0) { + throw std::runtime_error("SnakeBeta1dConfig.hidden_size must be positive"); + } +} + +const SnakeBeta1dConfig & SnakeBeta1dModule::config() const noexcept { + return config_; +} + +const core::ModuleSchema & SnakeBeta1dModule::schema() const noexcept { + return static_schema(); +} + +core::TensorValue SnakeBeta1dModule::build( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const SnakeBeta1dWeights & weights) const { + if (ctx.ggml == nullptr) { + throw std::runtime_error("ModuleBuildContext.ggml is null"); + } + core::validate_rank_between(input, 2, core::kMaxTensorRank, "input"); + if (input.shape.dims[input.shape.rank - 2] != config_.hidden_size) { + throw std::runtime_error("SnakeBeta1d input hidden dimension mismatch"); + } + core::validate_shape(weights.alpha, core::TensorShape::from_dims({config_.hidden_size}), "alpha"); + core::validate_shape(weights.beta, core::TensorShape::from_dims({config_.hidden_size}), "beta"); + + const auto contiguous = core::ensure_backend_addressable_layout(ctx, input); + const auto input_f32 = ensure_f32(ctx, contiguous); + const auto channel_shape = make_snake_alpha_shape(input.shape, config_.hidden_size); + auto alpha = core::reshape_tensor(ctx, ensure_f32(ctx, weights.alpha), channel_shape); + auto beta = core::reshape_tensor(ctx, ensure_f32(ctx, weights.beta), channel_shape); + if (config_.logscale) { + alpha = core::wrap_tensor(ggml_exp(ctx.ggml, alpha.tensor), alpha.shape, GGML_TYPE_F32); + beta = core::wrap_tensor(ggml_exp(ctx.ggml, beta.tensor), beta.shape, GGML_TYPE_F32); + } + const auto ax = core::wrap_tensor(ggml_mul(ctx.ggml, input_f32.tensor, alpha.tensor), input_f32.shape, GGML_TYPE_F32); + const auto s = core::wrap_tensor(ggml_sin(ctx.ggml, ax.tensor), input_f32.shape, GGML_TYPE_F32); + const auto s2 = core::wrap_tensor(ggml_mul(ctx.ggml, s.tensor, s.tensor), input_f32.shape, GGML_TYPE_F32); + const auto denom = core::wrap_tensor(ggml_scale_bias(ctx.ggml, beta.tensor, 1.0F, 1.0e-9F), beta.shape, GGML_TYPE_F32); + const auto periodic = core::wrap_tensor(ggml_div(ctx.ggml, s2.tensor, denom.tensor), input_f32.shape, GGML_TYPE_F32); + return core::wrap_tensor(ggml_add(ctx.ggml, input_f32.tensor, periodic.tensor), input_f32.shape, GGML_TYPE_F32); +} + +const core::ModuleSchema & SnakeBeta1dModule::static_schema() noexcept { + return kSnakeBeta1dSchema; +} + AliasFreeActivationModule::AliasFreeActivationModule(AliasFreeActivationConfig config) : config_(config) { if (config_.channels <= 0) { throw std::runtime_error("AliasFreeActivationConfig.channels must be positive"); diff --git a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp index 811985a6..67a703b6 100644 --- a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp +++ b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp @@ -116,6 +116,20 @@ QwenCausalDecoderWeights causal_decoder_weights(const QwenCausalDecodeRuntimeWei return out; } +core::TensorValue apply_readback_rounding( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const QwenCausalDecodeRuntimeConfig & config, + core::BackendType backend_type) { + if (!config.readback_round_type.has_value()) { + return input; + } + if (*config.readback_round_type == GGML_TYPE_BF16 && backend_type != core::BackendType::Metal) { + return core::wrap_tensor(ggml_round_bf16(ctx.ggml, input.tensor), input.shape, GGML_TYPE_F32); + } + return input; +} + void round_readback(std::vector & values, const QwenCausalDecodeRuntimeConfig & config) { if (!config.readback_round_type.has_value()) { return; @@ -247,6 +261,52 @@ void write_batched_cached_step_mask( ggml_backend_tensor_set(tensor, scratch.data(), 0, scratch.size() * sizeof(ggml_fp16_t)); } +void write_batched_cached_step_mask_variable( + const QwenCausalDecodeRuntimeConfig & config, + ggml_tensor * tensor, + std::vector & scratch, + int64_t batch_size, + int64_t mask_steps, + const std::vector & visible_prefix_steps, + const std::vector & current_slots, + const std::vector & positions) { + if (tensor == nullptr) { + throw std::runtime_error("QwenCausalDecodeRuntime variable batched cached mask requires a tensor"); + } + if (batch_size <= 0 || mask_steps <= 0 || + visible_prefix_steps.size() != static_cast(batch_size) || + current_slots.size() != static_cast(batch_size) || + positions.size() != static_cast(batch_size)) { + throw std::runtime_error("QwenCausalDecodeRuntime variable batched cached mask shape mismatch"); + } + const auto masked = ggml_fp32_to_fp16(-std::numeric_limits::infinity()); + const auto visible = ggml_fp32_to_fp16(0.0F); + const size_t row_size = static_cast(mask_steps); + const size_t total_size = static_cast(batch_size) * row_size; + if (scratch.size() != total_size) { + scratch.resize(total_size); + } + std::fill(scratch.begin(), scratch.end(), masked); + for (int64_t batch = 0; batch < batch_size; ++batch) { + const int64_t visible_steps = visible_prefix_steps[static_cast(batch)]; + const int64_t current_slot = + current_slots[static_cast(batch)] - static_cast(batch * mask_steps); + if (visible_steps < 0 || visible_steps > mask_steps || current_slot < 0 || current_slot >= mask_steps) { + throw std::runtime_error("QwenCausalDecodeRuntime variable batched cached mask row range is invalid"); + } + const size_t row_offset = static_cast(batch) * row_size; + int64_t begin = 0; + if (config.sliding_window > 0) { + begin = std::max(0, positions[static_cast(batch)] - config.sliding_window + 1); + } + for (int64_t i = begin; i < visible_steps; ++i) { + scratch[row_offset + static_cast(i)] = visible; + } + scratch[row_offset + static_cast(current_slot)] = visible; + } + ggml_backend_tensor_set(tensor, scratch.data(), 0, scratch.size() * sizeof(ggml_fp16_t)); +} + core::TensorValue compact_logits_readback( core::ModuleBuildContext & ctx, const QwenCausalDecodeRuntimeConfig & config, @@ -461,7 +521,7 @@ class QwenCausalDecodeRuntime::Impl { if (token_ids.empty()) { throw std::runtime_error("QwenCausalDecodeRuntime prefill requires tokens"); } - ensure_prefill_token_graph(static_cast(token_ids.size())); + ensure_prefill_token_graph(static_cast(token_ids.size()), false); ggml_backend_tensor_set( prefill_input_, token_ids.data(), @@ -470,6 +530,25 @@ class QwenCausalDecodeRuntime::Impl { return run_prefill(); } + QwenCausalPrefillIntoDecodeResult prefill_tokens_into_decode_cache( + const std::vector & token_ids, + int64_t required_cache_steps) { + if (token_ids.empty()) { + throw std::runtime_error("QwenCausalDecodeRuntime prefill requires tokens"); + } + if (required_cache_steps <= 0) { + throw std::runtime_error("QwenCausalDecodeRuntime decode requires positive cache capacity"); + } + ensure_decode_token_graph(required_cache_steps); + ensure_prefill_token_graph(static_cast(token_ids.size()), true); + ggml_backend_tensor_set( + prefill_input_, + token_ids.data(), + 0, + token_ids.size() * sizeof(int32_t)); + return run_prefill_into_decode_cache(); + } + QwenCausalPrefillResult prefill_embeddings(const std::vector & embeddings, int64_t steps) { if (steps <= 0) { throw std::runtime_error("QwenCausalDecodeRuntime prefill requires positive embedding steps"); @@ -609,6 +688,15 @@ class QwenCausalDecodeRuntime::Impl { return run_decode_step(); } + void decode_token_into(int32_t token, QwenCausalDecodeStepResult & out) { + ensure_decode_started(); + if (decode_input_kind_ != InputKind::Token) { + throw std::runtime_error("QwenCausalDecodeRuntime decode graph expects embeddings"); + } + ggml_backend_tensor_set(decode_input_, &token, 0, sizeof(int32_t)); + run_decode_step_into(out); + } + QwenCausalDecodeStepResult decode_embedding(const std::vector & embedding) { ensure_decode_started(); if (decode_input_kind_ != InputKind::Embedding) { @@ -627,7 +715,9 @@ class QwenCausalDecodeRuntime::Impl { if (required_cache_steps <= 0) { throw std::runtime_error("QwenCausalDecodeRuntime batched decode requires positive cache capacity"); } - ensure_batched_decode_token_graph(required_cache_steps, state.batch_size); + const bool variable_positions = + !state.current_end_by_batch.empty() || !state.valid_steps_by_batch.empty(); + ensure_batched_decode_token_graph(required_cache_steps, state.batch_size, variable_positions); batched_decode_cache_.import_state(state); } @@ -641,7 +731,9 @@ class QwenCausalDecodeRuntime::Impl { if (required_cache_steps <= 0) { throw std::runtime_error("QwenCausalDecodeRuntime batched decode requires positive cache capacity"); } - ensure_batched_decode_embedding_graph(required_cache_steps, state.batch_size); + const bool variable_positions = + !state.current_end_by_batch.empty() || !state.valid_steps_by_batch.empty(); + ensure_batched_decode_embedding_graph(required_cache_steps, state.batch_size, variable_positions); batched_decode_cache_.import_state(state); } @@ -700,13 +792,14 @@ class QwenCausalDecodeRuntime::Impl { Embedding, }; - void ensure_prefill_token_graph(int64_t steps) { - if (prefill_graph_ != nullptr && prefill_input_kind_ == InputKind::Token && prefill_steps_ == steps) { + void ensure_prefill_token_graph(int64_t steps, bool populate_decode_cache) { + if (prefill_graph_ != nullptr && prefill_input_kind_ == InputKind::Token && prefill_steps_ == steps && + prefill_populates_decode_cache_ == populate_decode_cache) { debug::trace_log_scalar(config_.trace_name + ".prefill.steps", steps); return; } release_prefill_graph(); - build_prefill_graph(InputKind::Token, steps); + build_prefill_graph(InputKind::Token, steps, populate_decode_cache); } void ensure_prefill_embedding_graph(int64_t steps) { @@ -715,10 +808,18 @@ class QwenCausalDecodeRuntime::Impl { return; } release_prefill_graph(); - build_prefill_graph(InputKind::Embedding, steps); + build_prefill_graph(InputKind::Embedding, steps, false); } - void build_prefill_graph(InputKind input_kind, int64_t steps) { + void build_prefill_graph(InputKind input_kind, int64_t steps, bool populate_decode_cache) { + if (populate_decode_cache) { + if (decode_graph_ == nullptr) { + throw std::runtime_error("QwenCausalDecodeRuntime prefill-to-decode requires an initialized decode graph"); + } + if (steps > decode_cache_steps_) { + throw std::runtime_error("QwenCausalDecodeRuntime prefill-to-decode exceeds decode cache capacity"); + } + } const auto build_start = Clock::now(); ggml_init_params params{config_.prefill_graph_arena_bytes, nullptr, true}; prefill_ctx_.reset(ggml_init(params)); @@ -750,22 +851,51 @@ class QwenCausalDecodeRuntime::Impl { GGML_TYPE_F16); auto decoder_out = build_causal_prefill(ctx, config_, x, positions, weights_, attention_mask); prefill_logits_readback_token_ids_ = make_logits_readback_token_ids(prefill_ctx_.get(), config_); - for (const auto & layer : decoder_out.state.layers) { + if (populate_decode_cache) { + prefill_graph_ = ggml_new_graph_custom(prefill_ctx_.get(), 65536, false); + } + for (size_t layer_index = 0; layer_index < decoder_out.state.layers.size(); ++layer_index) { + const auto & layer = decoder_out.state.layers[layer_index]; if (!layer.key.has_value() || !layer.value.has_value()) { throw std::runtime_error("QwenCausalDecodeRuntime prefill decoder did not return K/V state"); } - auto * key = ggml_cpy( - prefill_ctx_.get(), - layer.key->tensor, - ggml_dup_tensor(prefill_ctx_.get(), layer.key->tensor)); - auto * value = ggml_cpy( - prefill_ctx_.get(), - layer.value->tensor, - ggml_dup_tensor(prefill_ctx_.get(), layer.value->tensor)); - ggml_set_output(key); - ggml_set_output(value); - prefill_keys_.push_back(key); - prefill_values_.push_back(value); + if (populate_decode_cache) { + auto key = apply_readback_rounding(ctx, *layer.key, config_, backend_type_); + auto value = apply_readback_rounding(ctx, *layer.value, config_, backend_type_); + auto key_dest = runtime::view_transformer_kv_cache_steps( + ctx, + decode_cache_.key_tensor(layer_index), + 0, + steps, + config_.decoder.stack.num_key_value_heads, + config_.decoder.stack.head_dim, + "QwenCausalDecodeRuntime prefill key cache", + decode_cache_.key_tensor(layer_index).type); + auto value_dest = runtime::view_transformer_kv_cache_steps( + ctx, + decode_cache_.value_tensor(layer_index), + 0, + steps, + config_.decoder.stack.num_key_value_heads, + config_.decoder.stack.head_dim, + "QwenCausalDecodeRuntime prefill value cache", + decode_cache_.value_tensor(layer_index).type); + ggml_build_forward_expand(prefill_graph_, ggml_cpy(ctx.ggml, key.tensor, key_dest.tensor)); + ggml_build_forward_expand(prefill_graph_, ggml_cpy(ctx.ggml, value.tensor, value_dest.tensor)); + } else { + auto * key = ggml_cpy( + prefill_ctx_.get(), + layer.key->tensor, + ggml_dup_tensor(prefill_ctx_.get(), layer.key->tensor)); + auto * value = ggml_cpy( + prefill_ctx_.get(), + layer.value->tensor, + ggml_dup_tensor(prefill_ctx_.get(), layer.value->tensor)); + ggml_set_output(key); + ggml_set_output(value); + prefill_keys_.push_back(key); + prefill_values_.push_back(value); + } } if (config_.output_mode == QwenCausalDecodeOutputMode::Logits) { auto logits = decoder_out.logits; @@ -789,7 +919,9 @@ class QwenCausalDecodeRuntime::Impl { ggml_dup_tensor(prefill_ctx_.get(), decoder_out.hidden.tensor)); ggml_set_output(prefill_hidden_); } - prefill_graph_ = ggml_new_graph_custom(prefill_ctx_.get(), 65536, false); + if (!populate_decode_cache) { + prefill_graph_ = ggml_new_graph_custom(prefill_ctx_.get(), 65536, false); + } for (auto * key : prefill_keys_) { ggml_build_forward_expand(prefill_graph_, key); } @@ -825,6 +957,7 @@ class QwenCausalDecodeRuntime::Impl { } prefill_input_kind_ = input_kind; prefill_steps_ = steps; + prefill_populates_decode_cache_ = populate_decode_cache; debug::timing_log_scalar( config_.trace_name + ".prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); @@ -832,6 +965,9 @@ class QwenCausalDecodeRuntime::Impl { } QwenCausalPrefillResult run_prefill() { + if (prefill_populates_decode_cache_) { + throw std::runtime_error("QwenCausalDecodeRuntime prefill graph populates decode cache"); + } // The gallocr considers the persistent position/mask inputs dead after // their last read inside a compute and may hand their memory to other // tensors, so a cached prefill graph must be re-fed before recompute. @@ -881,6 +1017,47 @@ class QwenCausalDecodeRuntime::Impl { return out; } + QwenCausalPrefillIntoDecodeResult run_prefill_into_decode_cache() { + if (!prefill_populates_decode_cache_) { + throw std::runtime_error("QwenCausalDecodeRuntime prefill graph does not populate decode cache"); + } + ggml_backend_tensor_set( + prefill_positions_, + prefill_positions_values_.data(), + 0, + prefill_positions_values_.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + prefill_attention_mask_, + prefill_attention_mask_values_.data(), + 0, + prefill_attention_mask_values_.size() * sizeof(ggml_fp16_t)); + if (prefill_logits_readback_token_ids_ != nullptr) { + upload_logits_readback_token_ids(prefill_logits_readback_token_ids_, config_); + } + core::set_backend_threads(backend_, threads_); + const ggml_status status = core::compute_backend_graph(backend_, prefill_graph_); + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("QwenCausalDecodeRuntime prefill graph compute failed"); + } + QwenCausalPrefillIntoDecodeResult out; + if (prefill_logits_ != nullptr) { + out.logits.resize(static_cast(ggml_nelements(prefill_logits_))); + ggml_backend_tensor_get(prefill_logits_, out.logits.data(), 0, out.logits.size() * sizeof(float)); + } + if (prefill_hidden_ != nullptr) { + out.hidden.resize(static_cast(ggml_nelements(prefill_hidden_))); + ggml_backend_tensor_get(prefill_hidden_, out.hidden.data(), 0, out.hidden.size() * sizeof(float)); + round_readback(out.hidden, config_); + } + decode_cache_.retain_prefix(0); + decode_cache_.advance_after_direct_append(prefill_steps_); + begin_decode_mask(); + out.current_end = decode_cache_.current_end(); + out.valid_steps = decode_cache_.valid_steps(); + return out; + } + void ensure_batched_prefill_token_graph(int64_t batch_size, int64_t steps) { if (batched_prefill_graph_ != nullptr && batched_prefill_input_kind_ == InputKind::Token && batched_prefill_batch_size_ == batch_size && batched_prefill_steps_ == steps) { @@ -1195,25 +1372,31 @@ class QwenCausalDecodeRuntime::Impl { debug::trace_log_scalar(config_.trace_name + ".decode.cache_steps", cache_steps); } - void ensure_batched_decode_token_graph(int64_t cache_steps, int64_t batch_size) { + void ensure_batched_decode_token_graph(int64_t cache_steps, int64_t batch_size, bool variable_positions) { if (batched_decode_graph_ != nullptr && batched_decode_input_kind_ == InputKind::Token && - batched_decode_cache_steps_ >= cache_steps && batched_decode_batch_size_ == batch_size) { + batched_decode_cache_steps_ >= cache_steps && batched_decode_batch_size_ == batch_size && + batched_decode_variable_positions_ == variable_positions) { return; } release_batched_decode_graph(); - build_batched_decode_graph(InputKind::Token, batch_size, cache_steps); + build_batched_decode_graph(InputKind::Token, batch_size, cache_steps, variable_positions); } - void ensure_batched_decode_embedding_graph(int64_t cache_steps, int64_t batch_size) { + void ensure_batched_decode_embedding_graph(int64_t cache_steps, int64_t batch_size, bool variable_positions) { if (batched_decode_graph_ != nullptr && batched_decode_input_kind_ == InputKind::Embedding && - batched_decode_cache_steps_ >= cache_steps && batched_decode_batch_size_ == batch_size) { + batched_decode_cache_steps_ >= cache_steps && batched_decode_batch_size_ == batch_size && + batched_decode_variable_positions_ == variable_positions) { return; } release_batched_decode_graph(); - build_batched_decode_graph(InputKind::Embedding, batch_size, cache_steps); + build_batched_decode_graph(InputKind::Embedding, batch_size, cache_steps, variable_positions); } - void build_batched_decode_graph(InputKind input_kind, int64_t batch_size, int64_t cache_steps) { + void build_batched_decode_graph( + InputKind input_kind, + int64_t batch_size, + int64_t cache_steps, + bool variable_positions) { if (batch_size <= 0) { throw std::runtime_error("QwenCausalDecodeRuntime batched decode requires positive batch size"); } @@ -1240,10 +1423,13 @@ class QwenCausalDecodeRuntime::Impl { batched_decode_input_ = input.tensor; x = input; } - batched_decode_positions_ = ggml_new_tensor_1d(batched_decode_ctx_.get(), GGML_TYPE_I32, 1); + batched_decode_positions_ = ggml_new_tensor_1d( + batched_decode_ctx_.get(), + GGML_TYPE_I32, + variable_positions ? batch_size : 1); auto positions = core::wrap_tensor( batched_decode_positions_, - core::TensorShape::from_dims({1}), + core::TensorShape::from_dims({variable_positions ? batch_size : 1}), GGML_TYPE_I32); auto slot = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({batch_size})); batched_decode_cache_slot_ = slot.tensor; @@ -1308,6 +1494,7 @@ class QwenCausalDecodeRuntime::Impl { batched_decode_batch_size_ = batch_size; batched_decode_cache_steps_ = cache_steps; batched_decode_input_kind_ = input_kind; + batched_decode_variable_positions_ = variable_positions; debug::timing_log_scalar( config_.trace_name + ".batched_decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); @@ -1325,6 +1512,29 @@ class QwenCausalDecodeRuntime::Impl { } } + void begin_decode_mask() { + if (config_.sliding_window > 0) { + return; + } + if (decode_attention_mask_ == nullptr) { + throw std::runtime_error("QwenCausalDecodeRuntime decode attention mask is not initialized"); + } + std::fill( + decode_attention_mask_values_.begin(), + decode_attention_mask_values_.end(), + ggml_fp32_to_fp16(-std::numeric_limits::infinity())); + const int64_t visible_steps = std::min(decode_cache_.valid_steps(), decode_cache_steps_); + std::fill( + decode_attention_mask_values_.begin(), + decode_attention_mask_values_.begin() + static_cast(visible_steps), + ggml_fp32_to_fp16(0.0F)); + ggml_backend_tensor_set( + decode_attention_mask_, + decode_attention_mask_values_.data(), + 0, + decode_attention_mask_values_.size() * sizeof(ggml_fp16_t)); + } + QwenCausalDecodeStepResult run_decode_step() { if (decode_cache_.valid_steps() >= decode_cache_steps_) { throw std::runtime_error("QwenCausalDecodeRuntime decode cache exhausted"); @@ -1361,31 +1571,110 @@ class QwenCausalDecodeRuntime::Impl { return out; } + void run_decode_step_into(QwenCausalDecodeStepResult & out) { + if (decode_cache_.valid_steps() >= decode_cache_steps_) { + throw std::runtime_error("QwenCausalDecodeRuntime decode cache exhausted"); + } + const int32_t position = static_cast(decode_cache_.current_end()); + ggml_backend_tensor_set(decode_positions_, &position, 0, sizeof(int32_t)); + const int32_t cache_slot = static_cast(decode_cache_.valid_steps()); + ggml_backend_tensor_set(decode_cache_slot_, &cache_slot, 0, sizeof(int32_t)); + if (config_.sliding_window <= 0) { + const auto visible = ggml_fp32_to_fp16(0.0F); + decode_attention_mask_values_[static_cast(cache_slot)] = visible; + ggml_backend_tensor_set( + decode_attention_mask_, + &visible, + static_cast(cache_slot) * sizeof(ggml_fp16_t), + sizeof(ggml_fp16_t)); + } else { + write_cached_step_mask( + config_, + decode_attention_mask_, + decode_attention_mask_values_, + decode_cache_steps_, + decode_cache_.valid_steps(), + cache_slot, + position); + } + core::set_backend_threads(backend_, threads_); + const ggml_status status = core::compute_backend_graph(backend_, decode_graph_); + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("QwenCausalDecodeRuntime decode graph compute failed"); + } + if (decode_logits_ != nullptr) { + out.logits.resize(static_cast(ggml_nelements(decode_logits_))); + ggml_backend_tensor_get(decode_logits_, out.logits.data(), 0, out.logits.size() * sizeof(float)); + } else { + out.logits.clear(); + } + if (decode_hidden_ != nullptr) { + out.hidden.resize(static_cast(ggml_nelements(decode_hidden_))); + ggml_backend_tensor_get(decode_hidden_, out.hidden.data(), 0, out.hidden.size() * sizeof(float)); + round_readback(out.hidden, config_); + } else { + out.hidden.clear(); + } + decode_cache_.advance_after_direct_append(1); + } + QwenCausalDecodeStepResult run_batched_decode_step() { if (batched_decode_cache_.valid_steps() >= batched_decode_cache_steps_) { throw std::runtime_error("QwenCausalDecodeRuntime batched decode cache exhausted"); } - const int32_t position = static_cast(batched_decode_cache_.current_end()); - ggml_backend_tensor_set(batched_decode_positions_, &position, 0, sizeof(int32_t)); + int32_t position = static_cast(batched_decode_cache_.current_end()); const int32_t cache_slot = static_cast(batched_decode_cache_.valid_steps()); - for (int64_t batch = 0; batch < batched_decode_batch_size_; ++batch) { - batched_decode_cache_slots_[static_cast(batch)] = - static_cast(batch * batched_decode_cache_steps_ + cache_slot); + if (batched_decode_variable_positions_) { + const auto & current_ends = batched_decode_cache_.current_end_by_batch(); + const auto & valid_steps = batched_decode_cache_.valid_steps_by_batch(); + if (current_ends.size() != static_cast(batched_decode_batch_size_) || + valid_steps.size() != static_cast(batched_decode_batch_size_)) { + throw std::runtime_error("QwenCausalDecodeRuntime variable batched decode state is incomplete"); + } + batched_decode_positions_values_.resize(static_cast(batched_decode_batch_size_)); + for (int64_t batch = 0; batch < batched_decode_batch_size_; ++batch) { + const int64_t row_valid_steps = valid_steps[static_cast(batch)]; + batched_decode_positions_values_[static_cast(batch)] = + static_cast(current_ends[static_cast(batch)]); + batched_decode_cache_slots_[static_cast(batch)] = + static_cast(batch * batched_decode_cache_steps_ + row_valid_steps); + } + ggml_backend_tensor_set( + batched_decode_positions_, + batched_decode_positions_values_.data(), + 0, + batched_decode_positions_values_.size() * sizeof(int32_t)); + write_batched_cached_step_mask_variable( + config_, + batched_decode_attention_mask_, + batched_decode_attention_mask_values_, + batched_decode_batch_size_, + batched_decode_cache_steps_, + valid_steps, + batched_decode_cache_slots_, + current_ends); + } else { + ggml_backend_tensor_set(batched_decode_positions_, &position, 0, sizeof(int32_t)); + for (int64_t batch = 0; batch < batched_decode_batch_size_; ++batch) { + batched_decode_cache_slots_[static_cast(batch)] = + static_cast(batch * batched_decode_cache_steps_ + cache_slot); + } + write_batched_cached_step_mask( + config_, + batched_decode_attention_mask_, + batched_decode_attention_mask_values_, + batched_decode_batch_size_, + batched_decode_cache_steps_, + batched_decode_cache_.valid_steps(), + cache_slot, + position); } ggml_backend_tensor_set( batched_decode_cache_slot_, batched_decode_cache_slots_.data(), 0, batched_decode_cache_slots_.size() * sizeof(int32_t)); - write_batched_cached_step_mask( - config_, - batched_decode_attention_mask_, - batched_decode_attention_mask_values_, - batched_decode_batch_size_, - batched_decode_cache_steps_, - batched_decode_cache_.valid_steps(), - cache_slot, - position); core::set_backend_threads(backend_, threads_); const ggml_status status = core::compute_backend_graph(backend_, batched_decode_graph_); ggml_backend_synchronize(backend_); @@ -1437,6 +1726,7 @@ class QwenCausalDecodeRuntime::Impl { prefill_attention_mask_values_.clear(); prefill_steps_ = 0; prefill_input_kind_ = InputKind::None; + prefill_populates_decode_cache_ = false; } void release_batched_prefill_graph() { @@ -1530,6 +1820,9 @@ class QwenCausalDecodeRuntime::Impl { void release_decode_graph() { release_block_graph(); + if (prefill_populates_decode_cache_) { + release_prefill_graph(); + } if (decode_graph_ != nullptr) { core::release_backend_graph_resources( backend_, decode_graph_, config_.evict_cuda_graph_cache_on_release); @@ -1574,9 +1867,11 @@ class QwenCausalDecodeRuntime::Impl { batched_decode_cache_ = runtime::TransformerBatchedKVCache(); batched_decode_attention_mask_values_.clear(); batched_decode_cache_slots_.clear(); + batched_decode_positions_values_.clear(); batched_decode_batch_size_ = 0; batched_decode_cache_steps_ = 0; batched_decode_input_kind_ = InputKind::None; + batched_decode_variable_positions_ = false; } ggml_backend_t backend_ = nullptr; @@ -1603,6 +1898,7 @@ class QwenCausalDecodeRuntime::Impl { std::vector prefill_attention_mask_values_; int64_t prefill_steps_ = 0; InputKind prefill_input_kind_ = InputKind::None; + bool prefill_populates_decode_cache_ = false; std::unique_ptr batched_prefill_ctx_; ggml_tensor * batched_prefill_input_ = nullptr; @@ -1659,10 +1955,12 @@ class QwenCausalDecodeRuntime::Impl { ggml_backend_buffer_t batched_decode_buffer_ = nullptr; std::vector batched_decode_attention_mask_values_; std::vector batched_decode_cache_slots_; + std::vector batched_decode_positions_values_; runtime::TransformerBatchedKVCache batched_decode_cache_; int64_t batched_decode_batch_size_ = 0; int64_t batched_decode_cache_steps_ = 0; InputKind batched_decode_input_kind_ = InputKind::None; + bool batched_decode_variable_positions_ = false; }; QwenCausalDecodeRuntime::QwenCausalDecodeRuntime( @@ -1683,6 +1981,12 @@ QwenCausalPrefillResult QwenCausalDecodeRuntime::prefill_embeddings( return impl_->prefill_embeddings(embeddings, steps); } +QwenCausalPrefillIntoDecodeResult QwenCausalDecodeRuntime::prefill_tokens_into_decode_cache( + const std::vector & token_ids, + int64_t required_cache_steps) { + return impl_->prefill_tokens_into_decode_cache(token_ids, required_cache_steps); +} + QwenCausalBatchedPrefillResult QwenCausalDecodeRuntime::prefill_tokens_batched( const std::vector & token_ids, int64_t batch_size, @@ -1718,6 +2022,10 @@ QwenCausalDecodeStepResult QwenCausalDecodeRuntime::decode_token(int32_t token) return impl_->decode_token(token); } +void QwenCausalDecodeRuntime::decode_token_into(int32_t token, QwenCausalDecodeStepResult & out) { + impl_->decode_token_into(token, out); +} + QwenCausalDecodeStepResult QwenCausalDecodeRuntime::decode_embedding(const std::vector & embedding) { return impl_->decode_embedding(embedding); } diff --git a/src/framework/modules/transformers/qwen_decoder.cpp b/src/framework/modules/transformers/qwen_decoder.cpp index aa4b60c2..f5fba1d3 100644 --- a/src/framework/modules/transformers/qwen_decoder.cpp +++ b/src/framework/modules/transformers/qwen_decoder.cpp @@ -224,6 +224,44 @@ core::TensorValue cache_view( GGML_TYPE_F32); } +void apply_batched_static_rope( + core::ModuleBuildContext & ctx, + const QwenDecoderLayerConfig & config, + const QwenDecoderLayerWeights & weights, + const core::TensorValue & positions, + core::TensorValue & q, + core::TensorValue & k, + int64_t dim) { + const core::TensorValue * rope_factors = weights.rope_frequency_factors.has_value() + ? &*weights.rope_frequency_factors + : nullptr; + if (positions.shape.rank == 1 && positions.shape.dims[0] == q.shape.dims[1]) { + q = RoPEModule({dim, config.rope_type, config.rope_theta}).build(ctx, q, positions, rope_factors); + k = RoPEModule({dim, config.rope_type, config.rope_theta}).build(ctx, k, positions, rope_factors); + return; + } + if (q.shape.dims[1] != 1 || positions.shape.rank != 1 || positions.shape.dims[0] != q.shape.dims[0]) { + throw std::runtime_error("Qwen decoder batched static-cache RoPE positions must be [1] or [batch]"); + } + std::vector q_rows; + std::vector k_rows; + q_rows.reserve(static_cast(q.shape.dims[0])); + k_rows.reserve(static_cast(q.shape.dims[0])); + for (int64_t batch = 0; batch < q.shape.dims[0]; ++batch) { + auto q_row = SliceModule({0, batch, 1}).build(ctx, q); + auto k_row = SliceModule({0, batch, 1}).build(ctx, k); + auto pos_row = SliceModule({0, batch, 1}).build(ctx, positions); + if (ctx.backend_type == core::BackendType::Vulkan) { + // Vulkan RoPE cannot address a position view at a non-aligned byte offset. + pos_row = core::ensure_backend_addressable_layout(ctx, pos_row); + } + q_rows.push_back(RoPEModule({dim, config.rope_type, config.rope_theta}).build(ctx, q_row, pos_row, rope_factors)); + k_rows.push_back(RoPEModule({dim, config.rope_type, config.rope_theta}).build(ctx, k_row, pos_row, rope_factors)); + } + q = concat_all(ctx, q_rows, 0); + k = concat_all(ctx, k_rows, 0); +} + LinearWeights require_linear(const LinearWeights & weights, bool use_bias, const char * name) { if (use_bias && !weights.bias.has_value()) { throw std::runtime_error(std::string(name) + " bias is required"); @@ -921,11 +959,7 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail_bat auto v = reshape_qwen_heads(ctx, qkv.v, config_.num_key_value_heads, dim); if (config_.position_encoding == QwenDecoderPositionEncoding::Rotary) { - const core::TensorValue * rope_factors = weights.rope_frequency_factors.has_value() - ? &*weights.rope_frequency_factors - : nullptr; - q = RoPEModule({dim, config_.rope_type, config_.rope_theta}).build(ctx, q, positions, rope_factors); - k = RoPEModule({dim, config_.rope_type, config_.rope_theta}).build(ctx, k, positions, rope_factors); + apply_batched_static_rope(ctx, config_, weights, positions, q, k, dim); if (config_.activation_cast.enabled && config_.activation_cast.after_rope) { q = activation_cast(ctx, q, config_.activation_cast); k = activation_cast(ctx, k, config_.activation_cast); diff --git a/src/framework/runtime/kv_cache.cpp b/src/framework/runtime/kv_cache.cpp index 4a8450a0..0ee1000d 100644 --- a/src/framework/runtime/kv_cache.cpp +++ b/src/framework/runtime/kv_cache.cpp @@ -329,7 +329,19 @@ void TransformerBatchedKVCache::import_state(const TransformerBatchedKVState & s if (state.batch_size != batch_size_) { throw std::runtime_error("TransformerBatchedKVCache state batch size does not match cache batch size"); } - current_end_ = state.current_end; + if (!state.valid_steps_by_batch.empty() && + state.valid_steps_by_batch.size() != static_cast(batch_size_)) { + throw std::runtime_error("TransformerBatchedKVCache valid_steps_by_batch size mismatch"); + } + if (!state.current_end_by_batch.empty() && + state.current_end_by_batch.size() != static_cast(batch_size_)) { + throw std::runtime_error("TransformerBatchedKVCache current_end_by_batch size mismatch"); + } + current_end_by_batch_ = state.current_end_by_batch; + valid_steps_by_batch_ = state.valid_steps_by_batch; + current_end_ = current_end_by_batch_.empty() + ? state.current_end + : *std::max_element(current_end_by_batch_.begin(), current_end_by_batch_.end()); if (layers_.empty()) { valid_steps_ = 0; return; @@ -337,27 +349,44 @@ void TransformerBatchedKVCache::import_state(const TransformerBatchedKVState & s if (state.layers.size() != layers_.size()) { throw std::runtime_error("TransformerBatchedKVCache state layer count does not match cache layer count"); } - const int64_t state_steps = state.layers.empty() ? 0 : state.layers.front().valid_steps; + const int64_t state_steps = valid_steps_by_batch_.empty() + ? (state.layers.empty() ? 0 : state.layers.front().valid_steps) + : *std::max_element(valid_steps_by_batch_.begin(), valid_steps_by_batch_.end()); if (state_steps > cache_steps_) { throw std::runtime_error("TransformerBatchedKVCache state valid_steps exceeds cache capacity"); } valid_steps_ = state_steps; - const size_t copy_elems = static_cast(state_steps * row_elems_); for (size_t layer = 0; layer < layers_.size(); ++layer) { auto & cache = layers_[layer]; const auto & source = state.layers[layer]; - if (source.valid_steps != state_steps) { + if (valid_steps_by_batch_.empty() && source.valid_steps != state_steps) { throw std::runtime_error("TransformerBatchedKVCache requires consistent valid_steps across all layers"); } - const size_t state_elems = static_cast(batch_size_) * copy_elems; - if (source.key.size() != source.value.size() || source.key.size() != state_elems) { - throw std::runtime_error("TransformerBatchedKVCache source tensors do not match batch * valid_steps * row_elems"); - } std::fill(cache.import_key_scratch.begin(), cache.import_key_scratch.end(), 0.0F); std::fill(cache.import_value_scratch.begin(), cache.import_value_scratch.end(), 0.0F); + const bool variable_rows = !valid_steps_by_batch_.empty(); + size_t src_offset = 0; + if (!variable_rows) { + const size_t source_row_elems = static_cast(state_steps * row_elems_); + const size_t state_elems = static_cast(batch_size_) * source_row_elems; + if (source.key.size() != source.value.size() || source.key.size() != state_elems) { + throw std::runtime_error( + "TransformerBatchedKVCache source tensors do not match batch * valid_steps * row_elems"); + } + } for (int64_t batch = 0; batch < batch_size_; ++batch) { - const size_t src_offset = static_cast(batch) * copy_elems; + const int64_t row_steps = variable_rows ? valid_steps_by_batch_[static_cast(batch)] : state_steps; + if (row_steps < 0 || row_steps > state_steps) { + throw std::runtime_error("TransformerBatchedKVCache row valid_steps is invalid"); + } + const size_t copy_elems = static_cast(row_steps * row_elems_); const size_t dst_offset = static_cast(batch * cache_steps_ * row_elems_); + if (!variable_rows) { + src_offset = static_cast(batch) * static_cast(state_steps * row_elems_); + } + if (src_offset + copy_elems > source.key.size() || source.key.size() != source.value.size()) { + throw std::runtime_error("TransformerBatchedKVCache compact source tensor size mismatch"); + } std::copy( source.key.begin() + static_cast(src_offset), source.key.begin() + static_cast(src_offset + copy_elems), @@ -366,6 +395,12 @@ void TransformerBatchedKVCache::import_state(const TransformerBatchedKVState & s source.value.begin() + static_cast(src_offset), source.value.begin() + static_cast(src_offset + copy_elems), cache.import_value_scratch.begin() + static_cast(dst_offset)); + if (variable_rows) { + src_offset += copy_elems; + } + } + if (variable_rows && src_offset != source.key.size()) { + throw std::runtime_error("TransformerBatchedKVCache compact source tensor has trailing values"); } write_cache_tensor(cache.key_tensor, cache.import_key_scratch, options_); write_cache_tensor(cache.value_tensor, cache.import_value_scratch, options_); @@ -376,6 +411,8 @@ TransformerBatchedKVState TransformerBatchedKVCache::export_state() const { TransformerBatchedKVState state; state.batch_size = batch_size_; state.current_end = current_end_; + state.current_end_by_batch = current_end_by_batch_; + state.valid_steps_by_batch = valid_steps_by_batch_; state.layers.resize(layers_.size()); const size_t copy_elems = static_cast(valid_steps_ * row_elems_); const size_t state_elems = static_cast(batch_size_) * copy_elems; @@ -414,6 +451,12 @@ void TransformerBatchedKVCache::advance_after_direct_append(int64_t steps) { } valid_steps_ += steps; current_end_ += steps; + for (auto & value : valid_steps_by_batch_) { + value += steps; + } + for (auto & value : current_end_by_batch_) { + value += steps; + } } int64_t TransformerBatchedKVCache::batch_size() const noexcept { @@ -432,6 +475,14 @@ int64_t TransformerBatchedKVCache::cache_steps() const noexcept { return cache_steps_; } +const std::vector & TransformerBatchedKVCache::valid_steps_by_batch() const noexcept { + return valid_steps_by_batch_; +} + +const std::vector & TransformerBatchedKVCache::current_end_by_batch() const noexcept { + return current_end_by_batch_; +} + core::TensorValue view_transformer_kv_cache_steps( core::ModuleBuildContext & ctx, const core::TensorValue & cache, diff --git a/src/models/sheetsage/audio_frontend.cpp b/src/models/sheetsage/audio_frontend.cpp new file mode 100644 index 00000000..5e794c3e --- /dev/null +++ b/src/models/sheetsage/audio_frontend.cpp @@ -0,0 +1,162 @@ +#include "engine/models/sheetsage/audio_frontend.h" + +#include "engine/framework/audio/conversion.h" + +#include +#include +#include +#include +#include +#include + +namespace engine::models::sheetsage { + +std::vector SheetSage2AudioFrontend::prepare( + const std::vector & interleaved, + int source_rate, + int channels, + int target_rate, + int threads) { + if (source_rate <= 0 || target_rate <= 0 || channels < 1 || channels > 8 || threads < 1 || + interleaved.size() % static_cast(channels) != 0) { + throw std::runtime_error("SheetSage2 requires positive rates and 1-8 interleaved channels"); + } + // Default layouts: mono, stereo, 2.1, 4.0, 5.0, 5.1, 6.1, 7.1. + // Float mono output preserves front/surround mix gains without peak normalization. + constexpr float side = 0.7071067811865475244F; + constexpr std::array, 8> gains{{ + {{1}}, + {{side, side}}, + {{side, side, 0}}, + {{side, side, 1, 0.5F}}, + {{side, side, 1, 0.5F, 0.5F}}, + {{side, side, 1, 0, 0.5F, 0.5F}}, + {{side, side, 1, 0, 0.5F, 0.5F, 0.5F}}, + {{side, side, 1, 0, 0.5F, 0.5F, 0.5F, 0.5F}}, + }}; + const int64_t frames = static_cast(interleaved.size() / static_cast(channels)); + std::vector mono; + // Match the file frontend's operation order, including FP32 rounding. + if (channels > 1 && source_rate != target_rate) { + for (int channel = 0; channel < channels; ++channel) { + const auto resampled = prepare( + engine::audio::extract_interleaved_channel(interleaved, channels, channel), + source_rate, 1, target_rate, threads); + if (channel == 0) mono.resize(resampled.size(), 0.0F); + const float gain = gains[static_cast(channels - 1)][static_cast(channel)]; + for (size_t frame = 0; frame < resampled.size(); ++frame) { + mono[frame] += resampled[frame] * gain; + } + } + return mono; + } + if (channels == 1) { + mono = engine::audio::mixdown_interleaved_to_mono_average(interleaved, channels); + } else { + mono.resize(static_cast(frames)); + const auto & row = gains[static_cast(channels - 1)]; + for (int64_t frame = 0; frame < frames; ++frame) { + float value = 0.0F; + for (int channel = 0; channel < channels; ++channel) { + value += interleaved[static_cast(frame * channels + channel)] * row[static_cast(channel)]; + } + mono[static_cast(frame)] = value; + } + } + if (source_rate == target_rate || mono.empty()) { + return mono; + } + + // Python's file frontend uses a 32-lobe, beta=9 Kaiser sinc with 0.97 cutoff. + // Keep exact rational phases where possible, otherwise interpolate a 1024-phase bank. + constexpr double pi = 3.14159265358979323846264338327950288; + const double cutoff = std::min(0.97 * target_rate / source_rate, 1.0); + const int taps = (static_cast(std::ceil(32.0 / cutoff)) + 1) & ~1; + const int stride = (taps + 7) & ~7; + const int center = (taps - 1) / 2; + const int phases = std::min(1024, target_rate / std::gcd(source_rate, target_rate)); + const auto key = std::make_pair(source_rate, target_rate); + auto found = filters_.find(key); + if (found == filters_.end()) { + std::vector bank(static_cast((phases + 1) * stride)); + std::vector coefficients(static_cast(phases * taps)); + double normalization = 0.0; + for (int phase = 0; phase < phases; ++phase) { + for (int tap = 0; tap < taps; ++tap) { + const double distance = tap - center - static_cast(phase) / phases; + const double angle = pi * distance * cutoff; + const double radius = 2.0 * distance / taps; + const double x = 9.0 * std::sqrt(std::max(0.0, 1.0 - radius * radius)); + const double quarter_square = x * x / 4.0; + // I0(x) series for 0 <= x <= 9; the tail after k=32 is below 2e-31. + // Avoid special functions unavailable in Apple's libc++. + double window = 1.0; + double term = 1.0; + for (int k = 1; k <= 32; ++k) { + term *= quarter_square / (k * k); + window += term; + } + const double value = (angle == 0.0 ? 1.0 : std::sin(angle) / angle) * window; + if (phase == 0) { + normalization += value; + } + coefficients[static_cast(phase * taps + tap)] = value; + } + } + for (int phase = 0; phase < phases; ++phase) { + for (int tap = 0; tap < taps; ++tap) { + bank[static_cast(phase * stride + tap)] = static_cast( + coefficients[static_cast(phase * taps + tap)] / normalization); + } + } + // Phase 1 is phase 0 shifted by one input sample. + bank[static_cast(phases * stride)] = bank[static_cast(stride - 1)]; + std::copy_n(bank.begin(), stride - 1, bank.begin() + phases * stride + 1); + found = filters_.emplace(key, std::move(bank)).first; + } + const auto & bank = found->second; + // Flush only the reflected tail available after the unpadded convolution. + const int64_t unpadded_frames = std::max(0, + ((frames - taps / 2) * target_rate + source_rate - 1) / source_rate); + const int64_t remaining = frames - unpadded_frames * source_rate / target_rate + center; + const int64_t reflection = (std::min(remaining, taps) + 1) / 2; + const int64_t output_frames = std::max(0, + ((frames + reflection - taps / 2) * target_rate + source_rate - 1) / source_rate); + std::vector output(static_cast(output_frames)); +#ifdef _OPENMP +#pragma omp parallel for num_threads(threads) if (output_frames >= 4096) +#endif + for (int64_t frame = 0; frame < output_frames; ++frame) { + const int64_t numerator = frame * source_rate; + const int64_t sample = numerator / target_rate; + const int64_t phase_numerator = (numerator % target_rate) * phases; + const int phase = static_cast(phase_numerator / target_rate); + const double fraction = static_cast(phase_numerator % target_rate) / target_rate; + const float * weights = bank.data() + phase * stride; + // Fixed FP32 lanes keep the reduction independent of the inference backend. + std::array sums{}; + std::array next_sums{}; + for (int tap = 0; tap < stride; ++tap) { + int64_t index = sample - center + tap; + // Reflect about the first sample; repeat the last sample at the right edge. + while (index < 0 || index >= frames) { + index = index < 0 ? -index : 2 * frames - 1 - index; + } + const float input = mono[static_cast(index)]; + const size_t lane = static_cast(tap % 8); + sums[lane] = std::fma(input, weights[tap], sums[lane]); + if (fraction != 0.0) { + next_sums[lane] = std::fma(input, weights[stride + tap], next_sums[lane]); + } + } + const float value = ((sums[0] + sums[4]) + (sums[2] + sums[6])) + + ((sums[1] + sums[5]) + (sums[3] + sums[7])); + const float next_value = ((next_sums[0] + next_sums[4]) + (next_sums[2] + next_sums[6])) + + ((next_sums[1] + next_sums[5]) + (next_sums[3] + next_sums[7])); + output[static_cast(frame)] = static_cast( + fraction == 0.0 ? value : value + (next_value - value) * fraction); + } + return output; +} + +} // namespace engine::models::sheetsage diff --git a/src/models/sheetsage/processing.cpp b/src/models/sheetsage/processing.cpp new file mode 100644 index 00000000..0dcd9ad6 --- /dev/null +++ b/src/models/sheetsage/processing.cpp @@ -0,0 +1,943 @@ +#include "engine/models/sheetsage/processing.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::sheetsage { +namespace { + +constexpr std::array kStructureLabels = { + "silence", "intro", "outro", "verse", "chorus", "bridge", "pre-chorus", "post-chorus", + "interlude", "fade-out", "loop", "rap", "preshot", "irregular", "instrumental", + "intro and verse", "pre-chorus and chorus", "verse and pre-chorus", "solo", "theme", + "development", "variation", "pre-outro", +}; + +constexpr std::array kDurationTemplates = { + 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, + 96, 128, 192, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096, +}; + +constexpr std::array kChromaticSharps = { + "C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B", +}; + +const std::vector & full_chord_labels() { + static const std::vector labels = [] { + constexpr std::array qualities = { + "maj", "min", "dim", "aug", "maj7", "min7", "7", "hdim7", + "dim7", "minmaj7", "sus2", "sus4", "sus4(b7)", "maj6", "min6", + }; + const std::map> inversions = { + {"maj", {"/2", "/3", "/5"}}, + {"min", {"/2", "/b3", "/5"}}, + {"maj7", {"/3", "/5", "/7"}}, + {"min7", {"/b3", "/5", "/b7"}}, + {"7", {"/3", "/5", "/b7"}}, + }; + std::vector out; + out.reserve(361); + out.push_back("N"); + for (const auto * quality : qualities) { + const auto it = inversions.find(quality); + std::vector suffixes; + if (it != inversions.end()) { + suffixes = it->second; + } + suffixes.push_back(""); + for (const auto * root : kChromaticSharps) { + for (const auto * inversion : suffixes) { + out.push_back(std::string(root) + ":" + quality + inversion); + } + } + } + return out; + }(); + return labels; +} + +int key_accidental_count(const std::string & key) { + static const std::map values = { + {"C", 0}, {"G", 1}, {"D", 2}, {"A", 3}, {"E", 4}, {"B", 5}, {"F#", 6}, {"C#", 7}, + {"F", -1}, {"Bb", -2}, {"Eb", -3}, {"Ab", -4}, {"Db", -5}, {"Gb", -6}, {"Cb", -7}, + {"Am", 0}, {"Em", 1}, {"Bm", 2}, {"F#m", 3}, {"C#m", 4}, {"G#m", 5}, {"D#m", 6}, {"A#m", 7}, + {"Dm", -1}, {"Gm", -2}, {"Cm", -3}, {"Fm", -4}, {"Bbm", -5}, {"Ebm", -6}, {"Abm", -7}, + }; + const auto it = values.find(key); + return it == values.end() ? 0 : it->second; +} + +std::array key_accidentals(const std::string & key) { + std::array accidentals{}; + const int count = key_accidental_count(key); + const std::string order = count > 0 ? "FCGDAEB" : "BEADGCF"; + for (int i = 0; i < std::abs(count); ++i) { + const auto index = std::string("CDEFGAB").find(order[static_cast(i)]); + accidentals[index] = count > 0 ? 1 : -1; + } + return accidentals; +} + +std::string key_relative_pitch_name(int accidental_count, int pitch_class) { + static const std::map> names = { + {7, {"B#", "C#", "C##", "D#", "D##", "E#", "F#", "F##", "G#", "G##", "A#", "B"}}, + {6, {"B#", "C#", "C##", "D#", "E", "E#", "F#", "F##", "G#", "G##", "A#", "B"}}, + {5, {"B#", "C#", "C##", "D#", "E", "E#", "F#", "F##", "G#", "A", "A#", "B"}}, + {4, {"B#", "C#", "D", "D#", "E", "E#", "F#", "F##", "G#", "A", "A#", "B"}}, + {3, {"B#", "C#", "D", "D#", "E", "E#", "F#", "G", "G#", "A", "A#", "B"}}, + {2, {"C", "C#", "D", "D#", "E", "E#", "F#", "G", "G#", "A", "A#", "B"}}, + {1, {"C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"}}, + {0, {"C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "Bb", "B"}}, + {-1, {"C", "C#", "D", "Eb", "E", "F", "F#", "G", "G#", "A", "Bb", "B"}}, + {-2, {"C", "C#", "D", "Eb", "E", "F", "F#", "G", "Ab", "A", "Bb", "B"}}, + {-3, {"C", "Db", "D", "Eb", "E", "F", "F#", "G", "Ab", "A", "Bb", "B"}}, + {-4, {"C", "Db", "D", "Eb", "E", "F", "Gb", "G", "Ab", "A", "Bb", "B"}}, + {-5, {"C", "Db", "D", "Eb", "E", "F", "Gb", "G", "Ab", "A", "Bb", "Cb"}}, + {-6, {"C", "Db", "D", "Eb", "Fb", "F", "Gb", "G", "Ab", "A", "Bb", "Cb"}}, + {-7, {"C", "Db", "D", "Eb", "Fb", "F", "Gb", "G", "Ab", "Bbb", "Bb", "Cb"}}, + }; + return names.at(accidental_count)[static_cast(pitch_class)]; +} + +std::string note_to_abc( + int midi, + const std::array & key_accidentals, + std::map & measure_accidentals) { + const int accidental_count = std::accumulate(key_accidentals.begin(), key_accidentals.end(), 0); + const int pitch_class = ((midi % 12) + 12) % 12; + const std::string pitch_name = key_relative_pitch_name(accidental_count, pitch_class); + const char letter = pitch_name[0]; + const std::string accidental = pitch_name.substr(1); + const int accidental_number = + accidental == "bb" ? -2 : accidental == "b" ? -1 : accidental == "#" ? 1 : accidental == "##" ? 2 : 0; + int octave = (midi - 60) / 12; + if (pitch_class == 11 && accidental_number == -1) { + ++octave; + } else if (pitch_class == 0 && accidental_number == 1) { + --octave; + } + const int scale_index = static_cast(std::string("CDEFGAB").find(letter)); + const auto it = measure_accidentals.find(scale_index); + const int current = it == measure_accidentals.end() ? key_accidentals[static_cast(scale_index)] : it->second; + std::string prefix; + if (current != accidental_number) { + measure_accidentals[scale_index] = accidental_number; + prefix = accidental_number == -2 ? "__" : + accidental_number == -1 ? "_" : + accidental_number == 0 ? "=" : + accidental_number == 1 ? "^" : "^^"; + } + std::string note(1, letter); + if (octave > 0) { + note[0] = static_cast(std::tolower(static_cast(note[0]))); + note.append(static_cast(std::max(0, octave - 1)), '\''); + } else if (octave < 0) { + note.append(static_cast(-octave), ','); + } + return prefix + note; +} + +bool same_note_segment(int value, int next_value) { + if (value == 0) { + return next_value == 0; + } + const int pitch = value / 2 - 1; + return next_value == pitch * 2 + 2; +} + +bool continues_pitch(int value, int next_value) { + if (value <= 0) { + return false; + } + const int pitch = value / 2 - 1; + return next_value == pitch * 2 + 2; +} + +std::string abc_duration(int units) { + return units == 1 ? std::string{} : std::to_string(units); +} + +std::string key_symbol_to_abc(std::string key) { + const auto colon = key.find(':'); + if (colon != std::string::npos) { + auto root = key.substr(0, colon); + const auto mode = key.substr(colon + 1); + // The tokenizer uses sharp roots, including nonportable major keys. + if (mode == "major") { + static const std::map portable = { + {"A#", "Bb"}, {"D#", "Eb"}, {"G#", "Ab"}, + }; + const auto it = portable.find(root); + if (it != portable.end()) root = it->second; + } + return root + (mode == "minor" ? "m" : ""); + } + return key; +} + +std::string chord_symbol_to_abc(const std::string & chord) { + if (chord.empty() || chord == "N" || chord == "X" || chord == "?") { + return {}; + } + const auto colon = chord.find(':'); + if (colon == std::string::npos) { + return {}; + } + const auto root = chord.substr(0, colon); + auto quality = chord.substr(colon + 1); + const auto slash = quality.find('/'); + std::string bass; + if (slash != std::string::npos) { + bass = quality.substr(slash + 1); + quality = quality.substr(0, slash); + constexpr std::array natural = {0, 2, 4, 5, 7, 9, 11}; + const std::string letters = "CDEFGAB"; + const int root_letter = static_cast(letters.find(root[0])); + const int root_pitch = natural[static_cast(root_letter)] + + static_cast(std::count(root.begin(), root.end(), '#')) - + static_cast(std::count(root.begin(), root.end(), 'b')); + const auto degree_start = bass.find_first_of("123456789"); + const int degree = std::stoi(bass.substr(degree_start)); + const int alteration = static_cast(std::count(bass.begin(), bass.end(), '#')) - + static_cast(std::count(bass.begin(), bass.end(), 'b')); + const int letter = (root_letter + degree - 1) % 7; + const int pitch = (root_pitch + natural[static_cast((degree - 1) % 7)] + alteration + 12) % 12; + const int accidental = (pitch - natural[static_cast(letter)] + 18) % 12 - 6; + bass = std::string(1, letters[static_cast(letter)]) + + std::string(static_cast(std::abs(accidental)), accidental < 0 ? 'b' : '#'); + } + static const std::map quality_map = { + {"maj", ""}, {"min", "m"}, {"dim", "dim"}, {"aug", "aug"}, {"7", "7"}, + {"maj7", "maj7"}, {"min7", "m7"}, {"dim7", "dim7"}, {"hdim7", "m7b5"}, + {"sus4", "sus4"}, {"sus2", "sus2"}, {"maj6", "6"}, {"min6", "m6"}, + {"sus4(b7)", "7sus4"}, {"minmaj7", "m(maj7)"}, + }; + const auto it = quality_map.find(quality); + if (it == quality_map.end()) { + return {}; + } + return bass.empty() ? root + it->second : root + it->second + "/" + bass; +} + +struct BeatRow { + double time = 0.0; + int beat = 1; + int numerator = 4; + int denominator = 4; +}; + +struct NotationMeasure { + int start_beat = 0; + int end_beat = 0; + int numerator = 4; + int denominator = 4; + int abc_numerator = 4; + int abc_denominator = 4; + bool pad_before = false; +}; + +struct TextInterval { + double start = 0.0; + double end = 0.0; + std::string value; +}; + +std::vector rhythm_rows(const std::vector & events) { + std::vector rows; + std::optional> meter; + for (const auto & event : events) { + if (event.meter.has_value()) { + meter = event.meter; + } + if (!event.eighth_position.has_value() || !meter.has_value()) { + continue; + } + const int64_t scaled = *event.eighth_position * meter->second; + if (scaled % 8 != 0) { + continue; + } + const int beat = static_cast(scaled / 8) + 1; + if (beat < 1 || beat > meter->first) { + continue; + } + rows.push_back({event.time, beat, meter->first, meter->second}); + } + return rows; +} + +std::vector interval_rows( + const std::vector & events, + const char * field, + double duration) { + std::vector rows; + for (const auto & event : events) { + const std::optional * value = nullptr; + if (std::string(field) == "key") { + value = &event.key; + } else if (std::string(field) == "chord") { + value = &event.chord; + } else if (std::string(field) == "structure") { + value = &event.structure; + } + if (value != nullptr && value->has_value()) { + rows.push_back({event.time, duration, **value}); + } + } + for (size_t i = 0; i < rows.size(); ++i) { + rows[i].end = i + 1 < rows.size() ? rows[i + 1].start : duration; + } + rows.erase( + std::remove_if(rows.begin(), rows.end(), [](const auto & row) { return row.end <= row.start; }), + rows.end()); + return rows; +} + +std::vector infer_measures(const std::vector & beats) { + std::vector downbeats; + for (size_t i = 0; i < beats.size(); ++i) { + if (beats[i].beat == 1) { + downbeats.push_back(static_cast(i)); + } + } + if (downbeats.empty()) { + return {}; + } + std::vector measures; + const auto add_span = [&](int start, int end) { + if (end <= start) { + return; + } + std::map denom_counts; + std::map numerator_counts; + for (int i = start; i < end; ++i) { + ++denom_counts[beats[static_cast(i)].denominator]; + ++numerator_counts[beats[static_cast(i)].numerator]; + } + const auto best = [](const std::map & counts, int fallback) { + int value = fallback; + int count = -1; + for (const auto & [candidate, observed] : counts) { + if (observed > count) { + value = candidate; + count = observed; + } + } + return value; + }; + const int beat_count = end - start; + const int denominator = best(denom_counts, 4); + const int declared = best(numerator_counts, beat_count); + NotationMeasure measure; + measure.start_beat = start; + measure.end_beat = end; + measure.numerator = beat_count; + measure.denominator = denominator; + measure.abc_numerator = (end == static_cast(beats.size()) - 1 && declared >= beat_count) ? declared : beat_count; + measure.abc_denominator = denominator; + measures.push_back(measure); + }; + if (downbeats.front() > 0) { + add_span(0, downbeats.front()); + } + for (size_t i = 1; i < downbeats.size(); ++i) { + add_span(downbeats[i - 1], downbeats[i]); + } + if (downbeats.back() < static_cast(beats.size()) - 1) { + add_span(downbeats.back(), static_cast(beats.size()) - 1); + } + if (measures.size() >= 2) { + const double first = static_cast(measures.front().numerator) / measures.front().denominator; + const double following = static_cast(measures[1].abc_numerator) / measures[1].abc_denominator; + if (first < following) { + measures.front().abc_numerator = measures[1].abc_numerator; + measures.front().abc_denominator = measures[1].abc_denominator; + measures.front().pad_before = true; + } + } + return measures; +} + +std::vector subbeat_times_from_beats(const std::vector & beats) { + std::vector times; + if (beats.size() < 2) { + return times; + } + for (size_t i = 0; i + 1 < beats.size(); ++i) { + const double start = beats[i].time; + const double end = beats[i + 1].time; + for (int j = 0; j < 4; ++j) { + times.push_back(start + (end - start) * static_cast(j) / 4.0); + } + } + times.push_back(beats.back().time); + return times; +} + +int quantize_time(double time, const std::vector & subbeat_times) { + if (subbeat_times.size() < 2) { + return 0; + } + std::vector boundaries; + boundaries.reserve(subbeat_times.size() - 1); + for (size_t i = 0; i + 1 < subbeat_times.size(); ++i) { + boundaries.push_back((subbeat_times[i] + subbeat_times[i + 1]) * 0.5); + } + return static_cast(std::distance( + boundaries.begin(), + std::lower_bound(boundaries.begin(), boundaries.end(), time))); +} + +std::vector fill_text_intervals( + const std::vector & rows, + const std::vector & subbeat_times, + const std::string & fallback) { + std::vector out(subbeat_times.size(), fallback); + for (const auto & row : rows) { + int start = std::clamp(quantize_time(row.start, subbeat_times), 0, static_cast(out.size()) - 1); + int end = std::clamp(quantize_time(row.end, subbeat_times), 0, static_cast(out.size()) - 1); + if (end <= start) { + continue; + } + std::fill(out.begin() + start, out.begin() + end, row.value); + } + if (out.size() > 1) { + out.back() = out[out.size() - 2]; + } + return out; +} + +std::array, 2> build_voice_arrays( + const std::vector & events, + const std::vector & subbeat_times, + double duration) { + std::array, 2> voices = { + std::vector(subbeat_times.size(), 0), + std::vector(subbeat_times.size(), 0), + }; + struct TimedNote { + double start = 0.0; + double end = 0.0; + int pitch = 0; + int track = 0; + }; + std::vector notes; + for (const auto & event : events) { + for (size_t i = 0; i < event.notes.size(); ++i) { + const double end = i < event.note_end_times.size() + ? event.note_end_times[i] + : std::min(duration, static_cast(event.time) + 0.04); + notes.push_back({event.time, std::min(duration, end), event.notes[i].pitch, event.notes[i].track}); + } + } + std::sort(notes.begin(), notes.end(), [](const auto & a, const auto & b) { + if (a.track != b.track) { + return a.track < b.track; + } + if (a.start != b.start) { + return a.start < b.start; + } + if (a.pitch != b.pitch) { + return a.pitch < b.pitch; + } + return a.end < b.end; + }); + for (size_t i = 0; i < notes.size(); ++i) { + if (i + 1 < notes.size() && notes[i].track == notes[i + 1].track && notes[i].end > notes[i + 1].start) { + notes[i].end = notes[i + 1].start; + } + if (notes[i].end <= notes[i].start) { + continue; + } + const int track = std::clamp(notes[i].track, 0, 1); + int start = std::clamp(quantize_time(notes[i].start, subbeat_times), 0, static_cast(subbeat_times.size()) - 1); + int end = std::clamp(quantize_time(notes[i].end, subbeat_times), 0, static_cast(subbeat_times.size()) - 1); + if (end <= start) { + end = std::min(static_cast(subbeat_times.size()) - 1, start + 1); + } + const int sustain = notes[i].pitch * 2 + 2; + for (int t = start; t < end; ++t) { + if (voices[static_cast(track)][static_cast(t)] == 0) { + voices[static_cast(track)][static_cast(t)] = sustain; + } + } + voices[static_cast(track)][static_cast(start)] = sustain + 1; + } + return voices; +} + +std::vector subbeat_denominators_from_measures( + const std::vector & beats, + const std::vector & measures, + size_t subbeat_count) { + std::vector out(subbeat_count, 4); + for (const auto & measure : measures) { + for (int beat = measure.start_beat; beat < measure.end_beat; ++beat) { + for (int j = 0; j < 4; ++j) { + const size_t index = static_cast(beat * 4 + j); + if (index < out.size()) { + out[index] = measure.denominator; + } + } + } + } + if (out.size() > 1) { + out.back() = out[out.size() - 2]; + } + (void)beats; + return out; +} + +int duration_units( + const std::vector & denominators, + int start_t, + int end_t, + int unit_denominator) { + int units = 0; + for (int t = start_t; t < end_t; ++t) { + const int divisor = denominators[static_cast(t)] * 4; + units += unit_denominator / divisor; + } + return std::max(1, units); +} + +std::vector split_duration_units(int duration) { + static constexpr std::array supported = {48, 32, 24, 16, 12, 8, 6, 4, 3, 2, 1}; + std::vector out; + int remaining = std::max(1, duration); + while (remaining > 0) { + int chunk = 1; + for (const int candidate : supported) { + if (candidate <= remaining) { + chunk = candidate; + break; + } + } + out.push_back(chunk); + remaining -= chunk; + } + return out; +} + +std::vector render_duration_tokens( + const std::string & prefix, + const std::string & note, + int duration, + bool tie_out) { + const auto chunks = split_duration_units(duration); + std::vector out; + out.reserve(chunks.size()); + for (size_t i = 0; i < chunks.size(); ++i) { + const bool tie = note != "z" && (i + 1 < chunks.size() || tie_out); + out.push_back((i == 0 ? prefix : std::string{}) + note + abc_duration(chunks[i]) + (tie ? "-" : "")); + } + return out; +} + +std::string render_voice_measure( + const std::vector & voice, + const std::vector & keys, + const std::vector & chords, + const std::vector & denominators, + const NotationMeasure & measure, + int unit_denominator, + bool show_chords) { + std::ostringstream out; + int t = measure.start_beat * 4; + const int end_t = measure.end_beat * 4; + std::string current_key = keys[static_cast(t)]; + auto active_key_accidentals = key_accidentals(current_key); + std::map measure_accidentals; + int padding = measure.abc_numerator * unit_denominator / measure.abc_denominator - + measure.numerator * unit_denominator / measure.denominator; + int leading_padding = measure.pad_before ? padding : 0; + if (measure.pad_before) { + padding = 0; + } + while (t < end_t && t < static_cast(voice.size())) { + int next_t = std::min(end_t, static_cast(voice.size())); + for (int probe = t + 1; probe < next_t; ++probe) { + const bool note_change = !same_note_segment( + voice[static_cast(probe - 1)], + voice[static_cast(probe)]); + const bool attack = voice[static_cast(probe)] > 0 && voice[static_cast(probe)] % 2 == 1; + const bool key_change = keys[static_cast(probe)] != keys[static_cast(probe - 1)]; + const bool chord_change = show_chords && chords[static_cast(probe)] != chords[static_cast(probe - 1)]; + if (note_change || attack || key_change || chord_change) { + next_t = probe; + break; + } + } + std::string prefix; + if (t > measure.start_beat * 4 && keys[static_cast(t)] != keys[static_cast(t - 1)]) { + current_key = keys[static_cast(t)]; + active_key_accidentals = key_accidentals(current_key); + measure_accidentals.clear(); + prefix += "[K:" + current_key + "]"; + } + if (show_chords && (t == measure.start_beat * 4 || + chords[static_cast(t)] != chords[static_cast(t - 1)])) { + const auto chord = chord_symbol_to_abc(chords[static_cast(t)]); + if (!chord.empty()) { + prefix += "\"" + chord + "\""; + } + } + const int value = voice[static_cast(t)]; + const std::string note = value == 0 ? "z" : note_to_abc(value / 2 - 1, active_key_accidentals, measure_accidentals); + int units = duration_units(denominators, t, next_t, unit_denominator); + if (t == measure.start_beat * 4 && leading_padding > 0) { + if (value == 0 && prefix.empty()) { + units += leading_padding; + } else { + for (const auto & token : render_duration_tokens("", "z", leading_padding, false)) { + out << token; + } + } + leading_padding = 0; + } + if (value == 0 && next_t == end_t && padding > 0) { + units += padding; + padding = 0; + } + const bool tie_out = value > 0 && next_t < static_cast(voice.size()) && + continues_pitch(value, voice[static_cast(next_t)]); + for (const auto & token : render_duration_tokens(prefix, note, units, tie_out)) { + out << token; + } + t = next_t; + } + if (padding > 0) { + for (const auto & token : render_duration_tokens("", "z", padding, false)) { + out << token; + } + } + return out.str(); +} + +bool is_full_rest_measure(const std::string & text) { + return !text.empty() && text.find_first_not_of("z0123456789") == std::string::npos; +} + +std::string render_voice_group( + const std::vector & voice, + const std::vector & keys, + const std::vector & chords, + const std::vector & denominators, + const std::vector & measures, + int unit_denominator, + bool show_chords) { + std::ostringstream out; + std::vector rendered; + rendered.reserve(measures.size()); + for (const auto & measure : measures) { + rendered.push_back(render_voice_measure( + voice, + keys, + chords, + denominators, + measure, + unit_denominator, + show_chords)); + } + size_t index = 0; + while (index < rendered.size()) { + if (!is_full_rest_measure(rendered[index])) { + out << rendered[index] << "|"; + ++index; + continue; + } + size_t end = index + 1; + while (end < rendered.size() && is_full_rest_measure(rendered[end])) { + ++end; + } + const size_t count = end - index; + out << "Z"; + if (count > 1) { + out << count; + } + out << "|"; + index = end; + } + return out.str(); +} + +struct MeasureGroup { + std::vector measures; + std::vector structure_labels; + bool meter_changed = false; + bool key_changed = false; +}; + +std::string sanitize_structure_label(const std::string & value) { + std::istringstream in(value); + std::ostringstream out; + std::string part; + while (in >> part) { + if (out.tellp() > 0) { + out << " "; + } + out << part; + } + return out.str(); +} + +std::vector measure_groups( + const std::vector & measures, + const std::vector & keys, + const std::vector> & structure_events) { + if (measures.empty()) { + return {}; + } + std::pair active_meter{measures.front().abc_numerator, measures.front().abc_denominator}; + std::string active_key = keys[static_cast(measures.front().start_beat * 4)]; + std::string active_structure; + std::vector groups; + for (const auto & measure : measures) { + const auto meter = std::make_pair(measure.abc_numerator, measure.abc_denominator); + const std::string key = keys[static_cast(measure.start_beat * 4)]; + std::vector labels; + for (const auto & [t, label] : structure_events) { + if (t < measure.start_beat * 4 || t >= measure.end_beat * 4) { + continue; + } + const auto clean = sanitize_structure_label(label); + if (!clean.empty() && clean != active_structure) { + labels.push_back(clean); + active_structure = clean; + } + } + const bool meter_changed = meter != active_meter; + const bool key_changed = key != active_key; + const bool start_group = + groups.empty() || + groups.back().measures.size() >= 4 || + meter_changed || + key_changed || + !labels.empty(); + if (start_group) { + groups.push_back({{measure}, labels, meter_changed, key_changed}); + } else { + groups.back().measures.push_back(measure); + } + active_meter = meter; + active_key = keys[static_cast(std::max(0, measure.end_beat * 4 - 1))]; + } + return groups; +} + +} // namespace + +int64_t sheetsage2_structure_label_count() { + return static_cast(kStructureLabels.size()); +} + +int64_t sheetsage2_duration_bin_count() { + return static_cast(kDurationTemplates.size()); +} + +std::string sheetsage2_structure_label(int64_t index) { + index = std::clamp(index, 0, static_cast(kStructureLabels.size() - 1)); + return kStructureLabels[static_cast(index)]; +} + +std::string sheetsage2_key_label(int64_t index) { + return std::string(kChromaticSharps[static_cast(index % 12)]) + + (index >= 12 ? ":minor" : ":major"); +} + +std::string sheetsage2_chord_label(bool full_chord, int64_t index) { + if (full_chord) { + const auto & labels = full_chord_labels(); + index = std::clamp(index, 0, static_cast(labels.size() - 1)); + return labels[static_cast(index)]; + } + if (index == 0) { + return "N"; + } + const int64_t root = (index - 1) % 12; + const bool minor = (index - 1) >= 12; + return std::string(kChromaticSharps[static_cast(root)]) + (minor ? ":min" : ":maj"); +} + +SheetSage2Note sheetsage2_note_from_pitch_duration(int pitch_id, int64_t duration_bin) { + duration_bin = std::clamp(duration_bin, 0, static_cast(kDurationTemplates.size() - 1)); + return SheetSage2Note{ + pitch_id % 128, + pitch_id >= 128 ? 1 : 0, + static_cast(duration_bin), + kDurationTemplates[static_cast(duration_bin)], + }; +} + +std::string events_to_abc(const std::vector & events, double duration) { + auto beats = rhythm_rows(events); + if (beats.size() < 2) { + return "X:1\nT:\nM:4/4\nL:1/16\nQ:1/4=120\nV: Vocal clef=treble name=\"Vocal Melody\" snm=\"Vocal\"\nV: Ins clef=treble name=\"Ins Melody\" snm=\"Inst.\"\nK:C\nV: Vocal\nZ|\nV: Ins\nZ|\n"; + } + std::vector final_periods; + for (size_t i = beats.size() > 9 ? beats.size() - 8 : 1; i < beats.size(); ++i) { + final_periods.push_back(beats[i].time - beats[i - 1].time); + } + std::sort(final_periods.begin(), final_periods.end()); + const size_t middle = final_periods.size() / 2; + const double period = final_periods.size() % 2 == 0 + ? (final_periods[middle - 1] + final_periods[middle]) * 0.5 + : final_periods[middle]; + if (period <= 0.0) { + throw std::runtime_error("SheetSage2 decoded beats must increase in time"); + } + const double note_end = [&] { + double end = duration; + for (const auto & event : events) { + for (const auto value : event.note_end_times) { + end = std::max(end, static_cast(value)); + } + } + return end; + }(); + while (beats.back().time < note_end - 1.0e-6) { + auto next = beats.back(); + next.time += period; + next.beat = next.beat % next.numerator + 1; + beats.push_back(next); + } + auto measures = infer_measures(beats); + if (measures.empty()) { + return {}; + } + const auto subbeat_times = subbeat_times_from_beats(beats); + const auto denominators = subbeat_denominators_from_measures(beats, measures, subbeat_times.size()); + int unit_denominator = 1; + for (const auto & measure : measures) { + unit_denominator = std::lcm(unit_denominator, measure.denominator * 4); + unit_denominator = std::lcm(unit_denominator, measure.abc_denominator * 4); + } + const auto key_rows = interval_rows(events, "key", duration); + const auto chord_rows = interval_rows(events, "chord", duration); + const auto structure_rows = interval_rows(events, "structure", duration); + auto keys = fill_text_intervals( + key_rows, + subbeat_times, + key_rows.empty() ? std::string("C") : key_symbol_to_abc(key_rows.front().value)); + auto chords = fill_text_intervals(chord_rows, subbeat_times, "N"); + for (auto & key : keys) { + key = key_symbol_to_abc(key); + } + const auto voices = build_voice_arrays(events, subbeat_times, std::max(duration, note_end)); + const double seconds = subbeat_times.back() - subbeat_times.front(); + const double quarters = std::accumulate(denominators.begin(), denominators.end() - 1, 0.0, [](double sum, int denominator) { + return sum + 4.0 / static_cast(denominator) / 4.0; + }); + const int tempo = seconds > 0.0 ? static_cast(std::llround(quarters / seconds * 60.0)) : 120; + std::ostringstream out; + out << "X:1\n" + << "T:\n" + << "M:" << measures.front().abc_numerator << "/" << measures.front().abc_denominator << "\n" + << "L:1/" << unit_denominator << "\n" + << "Q:1/4=" << tempo << "\n" + << "V: Vocal clef=treble name=\"Vocal Melody\" snm=\"Vocal\"\n" + << "V: Ins clef=treble name=\"Ins Melody\" snm=\"Inst.\"\n" + << "K:" << keys[static_cast(measures.front().start_beat * 4)] << "\n"; + std::vector> structure_events; + structure_events.reserve(structure_rows.size()); + for (const auto & row : structure_rows) { + structure_events.push_back({quantize_time(row.start, subbeat_times), row.value}); + } + for (const auto & group : measure_groups(measures, keys, structure_events)) { + for (const auto & label : group.structure_labels) { + out << "% " << label << "\n"; + } + const auto & first = group.measures.front(); + out << "V: Vocal\n"; + if (group.meter_changed) { + out << "M:" << first.abc_numerator << "/" << first.abc_denominator << "\n"; + } + if (group.key_changed) { + out << "K:" << keys[static_cast(first.start_beat * 4)] << "\n"; + } + out << render_voice_group(voices[0], keys, chords, denominators, group.measures, unit_denominator, true) << "\n" + << "V: Ins\n"; + if (group.meter_changed) { + out << "M:" << first.abc_numerator << "/" << first.abc_denominator << "\n"; + } + if (group.key_changed) { + out << "K:" << keys[static_cast(first.start_beat * 4)] << "\n"; + } + out << render_voice_group(voices[1], keys, chords, denominators, group.measures, unit_denominator, false) << "\n"; + } + return out.str(); +} + +std::string events_json(const std::vector & events) { + const auto write_i32_array = [](std::ostringstream & out, const std::vector & values) { + out << "["; + for (size_t i = 0; i < values.size(); ++i) { + if (i != 0) { + out << ","; + } + out << values[i]; + } + out << "]"; + }; + std::ostringstream out; + out << "{\"events\":["; + for (size_t i = 0; i < events.size(); ++i) { + if (i != 0) { + out << ","; + } + const auto & event = events[i]; + out << "{\"subbeat\":" << event.subbeat + << ",\"time\":" << event.time + << ",\"window_index\":" << event.window_index + << ",\"source_subbeat\":" << event.source_subbeat + << ",\"global_subbeat\":" << event.global_subbeat + << ",\"tokens_by_field\":{"; + bool wrote = false; + const auto write_field = [&](const char * name, const std::vector & values) { + if (values.empty()) { + return; + } + if (wrote) { + out << ","; + } + wrote = true; + out << "\"" << name << "\":"; + write_i32_array(out, values); + }; + write_field("timestamp", event.timestamp_tokens); + write_field("rhythm", event.rhythm_tokens); + write_field("structure", event.structure_tokens); + write_field("key", event.key_tokens); + write_field("chord", event.chord_tokens); + write_field("melody", event.melody_tokens); + out << "},\"notes\":["; + for (size_t n = 0; n < event.notes.size(); ++n) { + if (n != 0) { + out << ","; + } + out << "{\"pitch\":" << event.notes[n].pitch + << ",\"track\":" << event.notes[n].track + << ",\"duration_bin\":" << event.notes[n].duration_bin + << ",\"duration_steps\":" << event.notes[n].duration_steps; + if (n < event.note_end_times.size()) { + out << ",\"end_time\":" << event.note_end_times[n]; + } + out << "}"; + } + out << "]}"; + } + out << "]}"; + return out.str(); +} + +} // namespace engine::models::sheetsage diff --git a/src/models/sheetsage/runtime.cpp b/src/models/sheetsage/runtime.cpp new file mode 100644 index 00000000..ab837b01 --- /dev/null +++ b/src/models/sheetsage/runtime.cpp @@ -0,0 +1,1552 @@ +#include "engine/models/sheetsage/runtime.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/attention/scaled_dot_product_attention.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/optimizations/fast_kv_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/runtime/kv_cache.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/streaming_conv_modules.h" +#include "engine/framework/modules/structural_modules.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::sheetsage { +namespace { + +namespace modules = engine::modules; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct SheetSage2AttentionWeights { + modules::LinearWeights q_proj; + modules::LinearWeights k_proj; + modules::LinearWeights v_proj; + modules::LinearWeights out_proj; +}; + +struct SheetSage2DecoderLayerWeights { + SheetSage2AttentionWeights self_attn; + modules::NormWeights self_attn_layer_norm; + SheetSage2AttentionWeights encoder_attn; + modules::NormWeights encoder_attn_layer_norm; + modules::LinearWeights fc1; + modules::LinearWeights fc2; + modules::NormWeights final_layer_norm; +}; + +struct SheetSage2DecoderWeights { + std::shared_ptr store; + core::TensorValue token_embedding; + core::TensorValue position_embedding; + modules::NormWeights layernorm_embedding; + modules::LinearWeights encoder_projection; + std::vector layers; +}; + +struct SheetSage2ConvNextLayerWeights { + modules::DepthwiseConv1dWeights depthwise; + modules::NormWeights norm; + modules::LinearWeights up; + modules::NormWeights grn; + modules::LinearWeights down; +}; + +struct SheetSage2SubsamplingBlockWeights { + std::optional resample_norm; + std::optional resample_conv; + std::vector layers; +}; + +struct SheetSage2EncoderLayerWeights { + modules::NormWeights ffn1_norm; + modules::LinearWeights ffn1_w1; + modules::LinearWeights ffn1_w2; + modules::NormWeights attn_norm; + SheetSage2AttentionWeights attn; + modules::NormWeights conv_norm; + modules::Conv1dWeights conv_pw_in; + modules::DepthwiseConv1dWeights conv_depthwise; + modules::NormWeights conv_depthwise_norm; + modules::Conv1dWeights conv_pw_out; + modules::NormWeights ffn2_norm; + modules::LinearWeights ffn2_w1; + modules::LinearWeights ffn2_w2; + modules::NormWeights final_norm; +}; + +struct SheetSage2EncoderWeights { + std::shared_ptr store; + core::TensorValue mel_mean; + core::TensorValue mel_std; + core::TensorValue half; + std::vector layer_weights; + std::vector mel_mean_host; + std::vector mel_std_host; + std::vector layer_weight_host; + std::vector subsampling; + std::vector layers; +}; + +void validate_config(const SheetSage2DecoderConfig & config) { + if (config.vocab_size <= 0 || config.hidden_size <= 0 || config.encoder_hidden_size <= 0 || + config.intermediate_size <= 0 || config.decoder_layers <= 0 || + config.num_attention_heads <= 0 || config.max_position_embeddings <= 0) { + throw std::runtime_error("SheetSage2 decoder config dimensions must be positive"); + } + if (config.hidden_size % config.num_attention_heads != 0) { + throw std::runtime_error("SheetSage2 decoder hidden size must be divisible by head count"); + } +} + +modules::LinearWeights load_linear( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t out_features, + int64_t in_features, + bool use_bias) { + modules::LinearWeights weights; + weights.weight = store.load_tensor(source, prefix + ".weight", storage_type, {out_features, in_features}); + if (use_bias) { + weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_features}); + } + return weights; +} + +modules::NormWeights load_norm( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t hidden) { + return { + store.load_f32_tensor(source, prefix + ".weight", {hidden}), + store.load_f32_tensor(source, prefix + ".bias", {hidden}), + }; +} + +SheetSage2AttentionWeights load_attention( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t hidden) { + return { + load_linear(store, source, prefix + ".q_proj", storage_type, hidden, hidden, true), + load_linear(store, source, prefix + ".k_proj", storage_type, hidden, hidden, true), + load_linear(store, source, prefix + ".v_proj", storage_type, hidden, hidden, true), + load_linear(store, source, prefix + ".out_proj", storage_type, hidden, hidden, true), + }; +} + +SheetSage2DecoderWeights load_weights( + const assets::TensorSource & source, + const SheetSage2DecoderConfig & config, + ggml_backend_t backend, + core::BackendType backend_type, + const SheetSage2DecoderRuntimeOptions & options) { + SheetSage2DecoderWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "models.sheetsage2.decoder.weights", + options.weight_context_bytes); + weights.token_embedding = weights.store->load_tensor( + source, + "token_embedding.weight", + options.weight_storage_type, + {config.vocab_size, config.hidden_size}); + weights.position_embedding = weights.store->load_tensor( + source, + "decoder.embed_positions.weight", + options.weight_storage_type, + {config.max_position_embeddings + 2, config.hidden_size}); + weights.layernorm_embedding = load_norm(*weights.store, source, "decoder.layernorm_embedding", config.hidden_size); + weights.encoder_projection = load_linear( + *weights.store, + source, + "encoder_projection", + options.weight_storage_type, + config.hidden_size, + config.encoder_hidden_size, + true); + weights.layers.reserve(static_cast(config.decoder_layers)); + for (int64_t i = 0; i < config.decoder_layers; ++i) { + const std::string prefix = "decoder.layers." + std::to_string(i); + SheetSage2DecoderLayerWeights layer; + layer.self_attn = load_attention(*weights.store, source, prefix + ".self_attn", options.weight_storage_type, config.hidden_size); + layer.self_attn_layer_norm = load_norm(*weights.store, source, prefix + ".self_attn_layer_norm", config.hidden_size); + layer.encoder_attn = load_attention(*weights.store, source, prefix + ".encoder_attn", options.weight_storage_type, config.hidden_size); + layer.encoder_attn_layer_norm = load_norm(*weights.store, source, prefix + ".encoder_attn_layer_norm", config.hidden_size); + layer.fc1 = load_linear(*weights.store, source, prefix + ".fc1", options.weight_storage_type, config.intermediate_size, config.hidden_size, true); + layer.fc2 = load_linear(*weights.store, source, prefix + ".fc2", options.weight_storage_type, config.hidden_size, config.intermediate_size, true); + layer.final_layer_norm = load_norm(*weights.store, source, prefix + ".final_layer_norm", config.hidden_size); + weights.layers.push_back(std::move(layer)); + } + weights.store->upload(); + return weights; +} + +modules::Conv1dWeights load_conv1d( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t out_channels, + int64_t in_channels, + int64_t kernel, + bool use_bias) { + modules::Conv1dWeights weights; + weights.weight = store.load_tensor(source, prefix + ".weight", storage_type, {out_channels, in_channels, kernel}); + if (use_bias) { + weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_channels}); + } + return weights; +} + +modules::DepthwiseConv1dWeights load_depthwise_conv1d( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t channels, + int64_t kernel, + bool use_bias) { + modules::DepthwiseConv1dWeights weights; + weights.weight = store.load_tensor(source, prefix + ".weight", storage_type, {channels, 1, kernel}); + if (use_bias) { + weights.bias = store.load_f32_tensor(source, prefix + ".bias", {channels}); + } + return weights; +} + +SheetSage2EncoderWeights load_encoder_weights( + const assets::TensorSource & source, + const SheetSage2DecoderConfig & config, + ggml_backend_t backend, + core::BackendType backend_type, + const SheetSage2DecoderRuntimeOptions & options) { + SheetSage2EncoderWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "models.sheetsage2.encoder.weights", + options.weight_context_bytes * 3); + weights.mel_mean_host = source.require_f32("feature_extractor.mel_mean", {config.mel_bins}); + weights.mel_std_host = source.require_f32("feature_extractor.mel_std", {config.mel_bins}); + weights.layer_weight_host = source.require_f32("layer_weight", {config.encoder_layers + 1}); + float max_weight = *std::max_element(weights.layer_weight_host.begin(), weights.layer_weight_host.end()); + float sum = 0.0F; + for (float & value : weights.layer_weight_host) { + value = std::exp(value - max_weight); + sum += value; + } + for (float & value : weights.layer_weight_host) { + value /= sum; + } + weights.mel_mean = weights.store->make_f32(core::TensorShape::from_dims({1, 1, config.mel_bins}), weights.mel_mean_host); + weights.mel_std = weights.store->make_f32(core::TensorShape::from_dims({1, 1, config.mel_bins}), weights.mel_std_host); + weights.half = weights.store->make_f32(core::TensorShape::from_dims({1, 1, 1}), {0.5F}); + weights.layer_weights.reserve(weights.layer_weight_host.size()); + for (const float value : weights.layer_weight_host) { + weights.layer_weights.push_back(weights.store->make_f32(core::TensorShape::from_dims({1, 1, 1}), {value})); + } + + const int64_t block_channels[] = {128, 512, 1024}; + const int64_t block_inputs[] = {128, 128, 512}; + const int64_t block_depths[] = {3, 4, 5}; + const int64_t block_strides[] = {1, 2, 2}; + weights.subsampling.reserve(3); + for (int64_t block = 0; block < 3; ++block) { + SheetSage2SubsamplingBlockWeights out; + const std::string prefix = "subsampling_module." + std::to_string(block); + if (block_inputs[block] != block_channels[block] || block_strides[block] > 1) { + out.resample_norm = load_norm(*weights.store, source, prefix + ".resampling_layer.0", block_inputs[block]); + out.resample_conv = load_conv1d( + *weights.store, + source, + prefix + ".resampling_layer.2", + options.weight_storage_type, + block_channels[block], + block_inputs[block], + 2, + true); + } + out.layers.reserve(static_cast(block_depths[block])); + for (int64_t layer = 0; layer < block_depths[block]; ++layer) { + const std::string layer_prefix = prefix + ".convnext_layers." + std::to_string(layer); + SheetSage2ConvNextLayerWeights cw; + cw.depthwise = load_depthwise_conv1d( + *weights.store, + source, + layer_prefix + ".depthwise_block.1", + options.weight_storage_type, + block_channels[block], + config.convnext_kernel_size, + true); + cw.norm = load_norm(*weights.store, source, layer_prefix + ".pointwise_block.0", block_channels[block]); + cw.up = load_linear( + *weights.store, + source, + layer_prefix + ".pointwise_block.1", + options.weight_storage_type, + block_channels[block] * 4, + block_channels[block], + true); + cw.grn = { + weights.store->load_f32_tensor(source, layer_prefix + ".pointwise_block.3.weight", {1, 1, block_channels[block] * 4}), + weights.store->load_f32_tensor(source, layer_prefix + ".pointwise_block.3.bias", {1, 1, block_channels[block] * 4}), + }; + cw.down = load_linear( + *weights.store, + source, + layer_prefix + ".pointwise_block.4", + options.weight_storage_type, + block_channels[block], + block_channels[block] * 4, + true); + out.layers.push_back(std::move(cw)); + } + weights.subsampling.push_back(std::move(out)); + } + weights.layers.reserve(static_cast(config.encoder_layers)); + for (int64_t i = 0; i < config.encoder_layers; ++i) { + const std::string prefix = "layers." + std::to_string(i); + SheetSage2EncoderLayerWeights layer; + layer.ffn1_norm = load_norm(*weights.store, source, prefix + ".ffn1_layer_norm", config.encoder_hidden_size); + layer.ffn1_w1 = load_linear(*weights.store, source, prefix + ".ffn1.w_1", options.weight_storage_type, config.encoder_intermediate_size, config.encoder_hidden_size, true); + layer.ffn1_w2 = load_linear(*weights.store, source, prefix + ".ffn1.w_2", options.weight_storage_type, config.encoder_hidden_size, config.encoder_intermediate_size, true); + layer.attn_norm = load_norm(*weights.store, source, prefix + ".attn_layer_norm", config.encoder_hidden_size); + layer.attn = { + load_linear(*weights.store, source, prefix + ".attn.query_proj", options.weight_storage_type, config.encoder_hidden_size, config.encoder_hidden_size, true), + load_linear(*weights.store, source, prefix + ".attn.key_proj", options.weight_storage_type, config.encoder_hidden_size, config.encoder_hidden_size, true), + load_linear(*weights.store, source, prefix + ".attn.value_proj", options.weight_storage_type, config.encoder_hidden_size, config.encoder_hidden_size, true), + load_linear(*weights.store, source, prefix + ".attn.out_proj", options.weight_storage_type, config.encoder_hidden_size, config.encoder_hidden_size, true), + }; + layer.conv_norm = load_norm(*weights.store, source, prefix + ".conv_module.layer_norm", config.encoder_hidden_size); + layer.conv_pw_in = load_conv1d(*weights.store, source, prefix + ".conv_module.conv_block.1", options.weight_storage_type, config.encoder_hidden_size * 2, config.encoder_hidden_size, 1, false); + layer.conv_depthwise = load_depthwise_conv1d(*weights.store, source, prefix + ".conv_module.conv_block.3", options.weight_storage_type, config.encoder_hidden_size, config.conformer_conv_kernel_size, false); + layer.conv_depthwise_norm = load_norm(*weights.store, source, prefix + ".conv_module.conv_block.4.1", config.encoder_hidden_size); + layer.conv_pw_out = load_conv1d(*weights.store, source, prefix + ".conv_module.conv_block.6", options.weight_storage_type, config.encoder_hidden_size, config.encoder_hidden_size, 1, false); + layer.ffn2_norm = load_norm(*weights.store, source, prefix + ".ffn2_layer_norm", config.encoder_hidden_size); + layer.ffn2_w1 = load_linear(*weights.store, source, prefix + ".ffn2.w_1", options.weight_storage_type, config.encoder_intermediate_size, config.encoder_hidden_size, true); + layer.ffn2_w2 = load_linear(*weights.store, source, prefix + ".ffn2.w_2", options.weight_storage_type, config.encoder_hidden_size, config.encoder_intermediate_size, true); + layer.final_norm = load_norm(*weights.store, source, prefix + ".final_layer_norm", config.encoder_hidden_size); + weights.layers.push_back(std::move(layer)); + } + weights.store->upload(); + return weights; +} + +core::TensorValue split_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t heads, + int64_t head_dim) { + auto shaped = core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, input), + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, head_dim})); + return modules::TransposeModule({{0, 2, 1, 3}, shaped.shape.rank}).build(ctx, shaped); +} + +core::TensorValue merge_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t hidden) { + return core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, input), + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], hidden})); +} + + +core::TensorValue attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & hidden, + const core::TensorValue & key_value, + const SheetSage2AttentionWeights & weights, + int64_t hidden_size, + int64_t heads, + bool causal) { + const int64_t head_dim = hidden_size / heads; + auto q = modules::LinearModule({hidden_size, hidden_size, true}).build(ctx, hidden, weights.q_proj); + auto k = modules::LinearModule({hidden_size, hidden_size, true}).build(ctx, key_value, weights.k_proj); + auto v = modules::LinearModule({hidden_size, hidden_size, true}).build(ctx, key_value, weights.v_proj); + q = split_heads(ctx, q, heads, head_dim); + k = split_heads(ctx, k, heads, head_dim); + v = split_heads(ctx, v, heads, head_dim); + auto context = modules::ScaledDotProductAttentionModule({ + head_dim, + modules::ScaledDotProductAttentionLowering::Explicit, + GGML_PREC_F32, + causal ? modules::AttentionCausality::Causal : modules::AttentionCausality::NonCausal, + }).build(ctx, q, k, v); + context = merge_heads(ctx, context, hidden_size); + return modules::LinearModule({hidden_size, hidden_size, true}).build(ctx, context, weights.out_proj); +} + +struct CachedSelfAttentionOutput { + core::TensorValue output; + core::TensorValue key_store; + core::TensorValue value_store; +}; + +CachedSelfAttentionOutput cached_self_attention_step( + core::ModuleBuildContext & ctx, + const core::TensorValue & hidden, + const core::TensorValue & cached_key_steps, + const core::TensorValue & cached_value_steps, + const core::TensorValue & cache_slot, + const core::TensorValue & attention_mask, + const SheetSage2AttentionWeights & weights, + int64_t hidden_size, + int64_t heads) { + const int64_t head_dim = hidden_size / heads; + auto q = modules::LinearModule({hidden_size, hidden_size, true}).build(ctx, hidden, weights.q_proj); + auto k = modules::LinearModule({hidden_size, hidden_size, true}).build(ctx, hidden, weights.k_proj); + auto v = modules::LinearModule({hidden_size, hidden_size, true}).build(ctx, hidden, weights.v_proj); + q = split_heads(ctx, q, heads, head_dim); + k = split_heads(ctx, k, heads, head_dim); + v = split_heads(ctx, v, heads, head_dim); + auto k_store = modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k); + auto v_store = modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v); + const modules::FastKVSetRowsModule set_rows({modules::FastKVSetRowsMode::BackendViewOptimized}); + auto updated_key_cache = set_rows.build(ctx, cached_key_steps, k_store, cache_slot); + auto updated_value_cache = set_rows.build(ctx, cached_value_steps, v_store, cache_slot); + auto all_k = modules::TransposeModule({{0, 2, 1, 3}, updated_key_cache.shape.rank}).build(ctx, updated_key_cache); + auto all_v = modules::TransposeModule({{0, 2, 1, 3}, updated_value_cache.shape.rank}).build(ctx, updated_value_cache); + auto context = modules::ScaledDotProductAttentionModule({ + head_dim, + modules::ScaledDotProductAttentionLowering::Explicit, + GGML_PREC_F32, + modules::AttentionCausality::NonCausal, + }).build(ctx, q, all_k, all_v, attention_mask); + context = merge_heads(ctx, context, hidden_size); + return { + modules::LinearModule({hidden_size, hidden_size, true}).build(ctx, context, weights.out_proj), + k_store, + v_store, + }; +} + +struct CrossAttentionKeyValue { + core::TensorValue key; + core::TensorValue value; +}; + +CrossAttentionKeyValue cross_attention_key_value( + core::ModuleBuildContext & ctx, + const core::TensorValue & memory, + const SheetSage2AttentionWeights & weights, + int64_t hidden_size, + int64_t heads) { + const int64_t head_dim = hidden_size / heads; + auto k = modules::LinearModule({hidden_size, hidden_size, true}).build(ctx, memory, weights.k_proj); + auto v = modules::LinearModule({hidden_size, hidden_size, true}).build(ctx, memory, weights.v_proj); + k = split_heads(ctx, k, heads, head_dim); + v = split_heads(ctx, v, heads, head_dim); + return {k, v}; +} + +core::TensorValue cached_cross_attention_step( + core::ModuleBuildContext & ctx, + const core::TensorValue & hidden, + const CrossAttentionKeyValue & key_value, + const SheetSage2AttentionWeights & weights, + int64_t hidden_size, + int64_t heads) { + const int64_t head_dim = hidden_size / heads; + auto q = modules::LinearModule({hidden_size, hidden_size, true}).build(ctx, hidden, weights.q_proj); + q = split_heads(ctx, q, heads, head_dim); + auto context = modules::ScaledDotProductAttentionModule({ + head_dim, + modules::ScaledDotProductAttentionLowering::Explicit, + GGML_PREC_F32, + modules::AttentionCausality::NonCausal, + }).build(ctx, q, key_value.key, key_value.value); + context = merge_heads(ctx, context, hidden_size); + return modules::LinearModule({hidden_size, hidden_size, true}).build(ctx, context, weights.out_proj); +} + +core::TensorValue decoder_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & memory, + const SheetSage2DecoderLayerWeights & weights, + const SheetSage2DecoderConfig & config) { + auto hidden = modules::ResidualAddModule{}.build( + ctx, + attention(ctx, input, input, weights.self_attn, config.hidden_size, config.num_attention_heads, true), + input); + hidden = modules::LayerNormModule({config.hidden_size, config.layer_norm_eps, true, true}).build( + ctx, + hidden, + weights.self_attn_layer_norm); + auto cross = attention(ctx, hidden, memory, weights.encoder_attn, config.hidden_size, config.num_attention_heads, false); + hidden = modules::ResidualAddModule{}.build(ctx, cross, hidden); + hidden = modules::LayerNormModule({config.hidden_size, config.layer_norm_eps, true, true}).build( + ctx, + hidden, + weights.encoder_attn_layer_norm); + auto ff = modules::LinearModule({config.hidden_size, config.intermediate_size, true}).build(ctx, hidden, weights.fc1); + ff = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, ff); + ff = modules::LinearModule({config.intermediate_size, config.hidden_size, true}).build(ctx, ff, weights.fc2); + hidden = modules::ResidualAddModule{}.build(ctx, ff, hidden); + return modules::LayerNormModule({config.hidden_size, config.layer_norm_eps, true, true}).build( + ctx, + hidden, + weights.final_layer_norm); +} + +struct CachedDecoderLayerOutput { + core::TensorValue hidden; + core::TensorValue key_store; + core::TensorValue value_store; +}; + +CachedDecoderLayerOutput cached_decoder_layer_step( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const CrossAttentionKeyValue & cross_key_value, + const core::TensorValue & cached_key_steps, + const core::TensorValue & cached_value_steps, + const core::TensorValue & cache_slot, + const core::TensorValue & attention_mask, + const SheetSage2DecoderLayerWeights & weights, + const SheetSage2DecoderConfig & config) { + auto self = cached_self_attention_step( + ctx, + input, + cached_key_steps, + cached_value_steps, + cache_slot, + attention_mask, + weights.self_attn, + config.hidden_size, + config.num_attention_heads); + auto hidden = modules::ResidualAddModule{}.build(ctx, self.output, input); + hidden = modules::LayerNormModule({config.hidden_size, config.layer_norm_eps, true, true}).build( + ctx, + hidden, + weights.self_attn_layer_norm); + auto cross = cached_cross_attention_step( + ctx, + hidden, + cross_key_value, + weights.encoder_attn, + config.hidden_size, + config.num_attention_heads); + hidden = modules::ResidualAddModule{}.build(ctx, cross, hidden); + hidden = modules::LayerNormModule({config.hidden_size, config.layer_norm_eps, true, true}).build( + ctx, + hidden, + weights.encoder_attn_layer_norm); + auto ff = modules::LinearModule({config.hidden_size, config.intermediate_size, true}).build(ctx, hidden, weights.fc1); + ff = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, ff); + ff = modules::LinearModule({config.intermediate_size, config.hidden_size, true}).build(ctx, ff, weights.fc2); + hidden = modules::ResidualAddModule{}.build(ctx, ff, hidden); + hidden = modules::LayerNormModule({config.hidden_size, config.layer_norm_eps, true, true}).build( + ctx, + hidden, + weights.final_layer_norm); + return {hidden, self.key_store, self.value_store}; +} + +core::TensorValue global_response_norm( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & eps, + const modules::NormWeights & weights) { + auto squared = core::wrap_tensor(ggml_sqr(ctx.ggml, input.tensor), input.shape, GGML_TYPE_F32); + auto sum = modules::ReduceSumModule({1}).build(ctx, squared); + auto magnitude = core::wrap_tensor(ggml_sqrt(ctx.ggml, sum.tensor), sum.shape, GGML_TYPE_F32); + auto mean = modules::ReduceMeanModule({2}).build(ctx, magnitude); + auto denom = modules::AddModule().build(ctx, mean, eps); + auto denom_full = modules::RepeatModule({magnitude.shape}).build(ctx, denom); + auto normalized = core::wrap_tensor(ggml_div(ctx.ggml, magnitude.tensor, denom_full.tensor), magnitude.shape, GGML_TYPE_F32); + auto normalized_full = modules::RepeatModule({input.shape}).build(ctx, normalized); + auto scaled = modules::MulModule().build(ctx, input, normalized_full); + auto weight = modules::RepeatModule({input.shape}).build(ctx, *weights.weight); + auto bias = modules::RepeatModule({input.shape}).build(ctx, *weights.bias); + auto out = modules::MulModule().build(ctx, scaled, weight); + out = modules::AddModule().build(ctx, out, bias); + return modules::AddModule().build(ctx, out, input); +} + +core::TensorValue convnext_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & grn_eps, + const SheetSage2ConvNextLayerWeights & weights, + int64_t channels, + float eps) { + auto bct = modules::TransposeModule({{0, 2, 1, 3}, input.shape.rank}).build(ctx, input); + bct = modules::DepthwiseConv1dModule({channels, 7, 1, 3, 1, true}).build(ctx, bct, weights.depthwise); + auto hidden = modules::TransposeModule({{0, 2, 1, 3}, bct.shape.rank}).build(ctx, bct); + hidden = modules::LayerNormModule({channels, eps, true, true}).build(ctx, hidden, weights.norm); + hidden = modules::LinearModule({channels, channels * 4, true}).build(ctx, hidden, weights.up); + hidden = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, hidden); + hidden = global_response_norm(ctx, hidden, grn_eps, weights.grn); + hidden = modules::LinearModule({channels * 4, channels, true}).build(ctx, hidden, weights.down); + return modules::ResidualAddModule().build(ctx, hidden, input); +} + +core::TensorValue subsampling_block( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & grn_eps, + const SheetSage2SubsamplingBlockWeights & weights, + int64_t in_channels, + int64_t out_channels, + int stride, + float eps) { + auto hidden = input; + if (weights.resample_norm.has_value() && weights.resample_conv.has_value()) { + hidden = modules::LayerNormModule({in_channels, eps, true, true}).build(ctx, hidden, *weights.resample_norm); + hidden = modules::TransposeModule({{0, 2, 1, 3}, hidden.shape.rank}).build(ctx, hidden); + hidden = modules::Conv1dModule({in_channels, out_channels, 2, stride, 0, 1, true}).build(ctx, hidden, *weights.resample_conv); + hidden = modules::TransposeModule({{0, 2, 1, 3}, hidden.shape.rank}).build(ctx, hidden); + } + for (const auto & layer : weights.layers) { + hidden = convnext_layer(ctx, hidden, grn_eps, layer, out_channels, eps); + } + return hidden; +} + +core::TensorValue encoder_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & hidden, + const core::TensorValue & cos, + const core::TensorValue & sin, + const SheetSage2AttentionWeights & weights, + const SheetSage2DecoderConfig & config) { + const int64_t head_dim = config.encoder_hidden_size / config.encoder_attention_heads; + auto q = modules::LinearModule({config.encoder_hidden_size, config.encoder_hidden_size, true}).build(ctx, hidden, weights.q_proj); + auto k = modules::LinearModule({config.encoder_hidden_size, config.encoder_hidden_size, true}).build(ctx, hidden, weights.k_proj); + auto v = modules::LinearModule({config.encoder_hidden_size, config.encoder_hidden_size, true}).build(ctx, hidden, weights.v_proj); + q = split_heads(ctx, q, config.encoder_attention_heads, head_dim); + k = split_heads(ctx, k, config.encoder_attention_heads, head_dim); + v = split_heads(ctx, v, config.encoder_attention_heads, head_dim); + q = modules::SplitRoPEModule({head_dim}).build(ctx, q, cos, sin); + k = modules::SplitRoPEModule({head_dim}).build(ctx, k, cos, sin); + auto context = modules::ScaledDotProductAttentionModule({ + head_dim, + modules::ScaledDotProductAttentionLowering::Explicit, + GGML_PREC_F32, + modules::AttentionCausality::NonCausal, + }).build(ctx, q, k, v); + context = merge_heads(ctx, context, config.encoder_hidden_size); + return modules::LinearModule({config.encoder_hidden_size, config.encoder_hidden_size, true}).build(ctx, context, weights.out_proj); +} + +core::TensorValue feed_forward( + core::ModuleBuildContext & ctx, + const core::TensorValue & hidden, + const modules::LinearWeights & w1, + const modules::LinearWeights & w2, + const SheetSage2DecoderConfig & config) { + auto out = modules::LinearModule({config.encoder_hidden_size, config.encoder_intermediate_size, true}).build(ctx, hidden, w1); + out = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, out); + return modules::LinearModule({config.encoder_intermediate_size, config.encoder_hidden_size, true}).build(ctx, out, w2); +} + +core::TensorValue conformer_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & cos, + const core::TensorValue & sin, + const core::TensorValue & half, + const SheetSage2EncoderLayerWeights & weights, + const SheetSage2DecoderConfig & config) { + auto hidden = modules::LayerNormModule({config.encoder_hidden_size, config.encoder_layer_norm_eps, true, true}).build(ctx, input, weights.ffn1_norm); + hidden = feed_forward(ctx, hidden, weights.ffn1_w1, weights.ffn1_w2, config); + auto out = modules::ResidualAddModule().build( + ctx, + modules::MulModule().build(ctx, hidden, modules::RepeatModule({hidden.shape}).build(ctx, half)), + input); + hidden = modules::LayerNormModule({config.encoder_hidden_size, config.encoder_layer_norm_eps, true, true}).build(ctx, out, weights.attn_norm); + hidden = encoder_attention(ctx, hidden, cos, sin, weights.attn, config); + out = modules::ResidualAddModule().build(ctx, hidden, out); + hidden = modules::LayerNormModule({config.encoder_hidden_size, config.encoder_layer_norm_eps, true, true}).build(ctx, out, weights.conv_norm); + hidden = modules::TransposeModule({{0, 2, 1, 3}, hidden.shape.rank}).build(ctx, hidden); + hidden = modules::Conv1dModule({config.encoder_hidden_size, config.encoder_hidden_size * 2, 1, 1, 0, 1, false}).build(ctx, hidden, weights.conv_pw_in); + auto gate_a = modules::SliceModule({1, 0, config.encoder_hidden_size}).build(ctx, hidden); + auto gate_b = modules::SliceModule({1, config.encoder_hidden_size, config.encoder_hidden_size}).build(ctx, hidden); + gate_b = core::wrap_tensor(ggml_sigmoid(ctx.ggml, gate_b.tensor), gate_b.shape, GGML_TYPE_F32); + hidden = modules::MulModule().build(ctx, gate_a, gate_b); + hidden = modules::DepthwiseConv1dModule({config.encoder_hidden_size, config.conformer_conv_kernel_size, 1, static_cast((config.conformer_conv_kernel_size - 1) / 2), 1, false}).build(ctx, hidden, weights.conv_depthwise); + hidden = modules::TransposeModule({{0, 2, 1, 3}, hidden.shape.rank}).build(ctx, hidden); + hidden = modules::LayerNormModule({config.encoder_hidden_size, config.encoder_layer_norm_eps, true, true}).build(ctx, hidden, weights.conv_depthwise_norm); + hidden = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, hidden); + hidden = modules::TransposeModule({{0, 2, 1, 3}, hidden.shape.rank}).build(ctx, hidden); + hidden = modules::Conv1dModule({config.encoder_hidden_size, config.encoder_hidden_size, 1, 1, 0, 1, false}).build(ctx, hidden, weights.conv_pw_out); + hidden = modules::TransposeModule({{0, 2, 1, 3}, hidden.shape.rank}).build(ctx, hidden); + out = modules::ResidualAddModule().build(ctx, hidden, out); + hidden = modules::LayerNormModule({config.encoder_hidden_size, config.encoder_layer_norm_eps, true, true}).build(ctx, out, weights.ffn2_norm); + hidden = feed_forward(ctx, hidden, weights.ffn2_w1, weights.ffn2_w2, config); + out = modules::ResidualAddModule().build( + ctx, + modules::MulModule().build(ctx, hidden, modules::RepeatModule({hidden.shape}).build(ctx, half)), + out); + return modules::LayerNormModule({config.encoder_hidden_size, config.encoder_layer_norm_eps, true, true}).build(ctx, out, weights.final_norm); +} + +} // namespace + +struct Mert2EncoderRuntime::Impl { + class EncoderGraph; + + Impl( + std::shared_ptr source, + core::ExecutionContext & execution, + SheetSage2DecoderConfig config, + SheetSage2DecoderRuntimeOptions options) + : source(std::move(source)), + execution(&execution), + config(config), + options(options) { + if (!this->source) { + throw std::runtime_error("MERT2 encoder runtime requires tensor source"); + } + validate_config(this->config); + } + + const SheetSage2EncoderWeights & require_encoder_weights() { + if (!encoder_weights) { + encoder_weights = std::make_unique(load_encoder_weights( + *source, + config, + execution->backend(), + execution->backend_type(), + options)); + source->release_storage(); + } + return *encoder_weights; + } + + std::shared_ptr source; + core::ExecutionContext * execution = nullptr; + SheetSage2DecoderConfig config; + SheetSage2DecoderRuntimeOptions options; + std::unique_ptr encoder_weights; + std::unique_ptr encoder_graph; +}; + +class Mert2EncoderRuntime::Impl::EncoderGraph { +public: + EncoderGraph( + core::ExecutionContext & execution, + const SheetSage2DecoderConfig & config, + const SheetSage2DecoderRuntimeOptions & options, + const SheetSage2EncoderWeights & weights, + int64_t batch, + int64_t mel_frames) + : backend_(execution.backend()), + backend_type_(execution.backend_type()), + threads_(std::max(1, execution.config().threads)), + config_(config), + options_(options), + weights_(weights), + batch_(batch), + mel_frames_(mel_frames) { + if (backend_ == nullptr || batch_ != 1 || mel_frames_ <= 0) { + throw std::runtime_error("SheetSage2 encoder graph initialization failed"); + } + encoded_frames_ = ((mel_frames_ - 2) / 2 + 1 - 2) / 2 + 1; + if (encoded_frames_ <= 0) { + throw std::runtime_error("SheetSage2 encoder graph computed no frames"); + } + build(); + } + + ~EncoderGraph() { + if (backend_ != nullptr && graph_ != nullptr) { + core::release_backend_graph_resources(backend_type_, backend_, graph_); + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + } + + bool matches(int64_t batch, int64_t mel_frames) const noexcept { + return batch == batch_ && mel_frames == mel_frames_; + } + + int64_t encoded_frames() const noexcept { + return encoded_frames_; + } + + std::vector run(const std::vector & normalized_mel) const { + if (static_cast(normalized_mel.size()) != batch_ * mel_frames_ * config_.mel_bins) { + throw std::runtime_error("SheetSage2 normalized mel shape mismatch"); + } + core::write_tensor_f32(mel_, normalized_mel); + core::write_tensor_f32(rope_cos_, rope_values(true)); + core::write_tensor_f32(rope_sin_, rope_values(false)); + core::write_tensor_f32(grn_eps_, std::vector(static_cast(batch_), 1.0e-6F)); + core::set_backend_threads(backend_, threads_); + const ggml_status status = core::compute_backend_graph(backend_, graph_, nullptr, "models.sheetsage2.encoder"); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("SheetSage2 encoder graph compute failed"); + } + return core::read_tensor_f32(mixed_); + } + +private: + std::vector rope_values(bool cosine) const { + const int64_t head_dim = config_.encoder_hidden_size / config_.encoder_attention_heads; + const int64_t half_dim = head_dim / 2; + std::vector values(static_cast(config_.encoder_attention_heads * encoded_frames_ * half_dim)); + for (int64_t h = 0; h < config_.encoder_attention_heads; ++h) { + for (int64_t t = 0; t < encoded_frames_; ++t) { + for (int64_t i = 0; i < half_dim; ++i) { + const float inv = std::pow(config_.rotary_embedding_base, -static_cast(2 * i) / static_cast(head_dim)); + const float value = cosine ? std::cos(static_cast(t) * inv) : std::sin(static_cast(t) * inv); + values[static_cast((h * encoded_frames_ + t) * half_dim + i)] = value; + } + } + } + return values; + } + + void build() { + ggml_init_params params{options_.graph_arena_bytes * 8, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("SheetSage2 encoder ggml context initialization failed"); + } + core::ModuleBuildContext input_ctx{ctx_.get(), "models.sheetsage2.encoder.inputs", backend_type_}; + mel_ = core::make_tensor( + input_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({batch_, mel_frames_, config_.mel_bins})); + const int64_t head_dim = config_.encoder_hidden_size / config_.encoder_attention_heads; + rope_cos_ = core::make_tensor( + input_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, config_.encoder_attention_heads, encoded_frames_, head_dim / 2})); + rope_sin_ = core::make_tensor( + input_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, config_.encoder_attention_heads, encoded_frames_, head_dim / 2})); + grn_eps_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_, 1, 1})); + ggml_set_input(mel_.tensor); + ggml_set_input(rope_cos_.tensor); + ggml_set_input(rope_sin_.tensor); + ggml_set_input(grn_eps_.tensor); + core::ModuleBuildContext build_ctx{ctx_.get(), "models.sheetsage2.encoder", backend_type_}; + auto mixed = build_graph_output(build_ctx); + mixed_ = mixed.tensor; + ggml_set_output(mixed_); + graph_ = ggml_new_graph_custom(ctx_.get(), 1048576, false); + ggml_build_forward_expand(graph_, mixed_); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (gallocr_ == nullptr || !ggml_gallocr_reserve(gallocr_, graph_) || !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + throw std::runtime_error("SheetSage2 encoder backend buffer allocation failed"); + } + } + + core::TensorValue build_graph_output(core::ModuleBuildContext & ctx) const { + auto hidden = mel_; + hidden = subsampling_block(ctx, hidden, grn_eps_, weights_.subsampling[0], 128, 128, 1, config_.subsampling_layer_norm_eps); + hidden = subsampling_block(ctx, hidden, grn_eps_, weights_.subsampling[1], 128, 512, 2, config_.subsampling_layer_norm_eps); + hidden = subsampling_block(ctx, hidden, grn_eps_, weights_.subsampling[2], 512, 1024, 2, config_.subsampling_layer_norm_eps); + auto mixed = modules::MulModule().build( + ctx, + hidden, + modules::RepeatModule({hidden.shape}).build(ctx, weights_.layer_weights[0])); + for (int64_t i = 0; i < config_.encoder_layers; ++i) { + hidden = conformer_layer( + ctx, + hidden, + rope_cos_, + rope_sin_, + weights_.half, + weights_.layers[static_cast(i)], + config_); + const auto layer_weight = modules::RepeatModule({hidden.shape}).build( + ctx, + weights_.layer_weights[static_cast(i + 1)]); + mixed = modules::AddModule().build(ctx, mixed, modules::MulModule().build(ctx, hidden, layer_weight)); + } + return mixed; + } + + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + SheetSage2DecoderConfig config_; + SheetSage2DecoderRuntimeOptions options_; + const SheetSage2EncoderWeights & weights_; + int64_t batch_ = 1; + int64_t mel_frames_ = 0; + int64_t encoded_frames_ = 0; + std::unique_ptr ctx_; + core::TensorValue mel_; + core::TensorValue rope_cos_; + core::TensorValue rope_sin_; + core::TensorValue grn_eps_; + ggml_tensor * mixed_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; +}; + +struct SheetSage2DecoderRuntime::Impl { + class DecodeGraph; + class CachedDecodeGraph; + + Impl( + std::shared_ptr source, + core::ExecutionContext & execution, + SheetSage2DecoderConfig config, + SheetSage2DecoderRuntimeOptions options) + : source(std::move(source)), + execution(&execution), + config(config), + options(options) { + if (!this->source) { + throw std::runtime_error("SheetSage2 decoder runtime requires tensor source"); + } + validate_config(this->config); + } + + const SheetSage2DecoderWeights & require_weights() { + if (!weights) { + weights = std::make_unique(load_weights( + *source, + config, + execution->backend(), + execution->backend_type(), + options)); + source->release_storage(); + } + return *weights; + } + + std::shared_ptr source; + core::ExecutionContext * execution = nullptr; + SheetSage2DecoderConfig config; + SheetSage2DecoderRuntimeOptions options; + std::unique_ptr weights; + std::unique_ptr graph; + std::unique_ptr cached_graph; +}; + +class SheetSage2DecoderRuntime::Impl::DecodeGraph { +public: + DecodeGraph( + core::ExecutionContext & execution, + const SheetSage2DecoderConfig & config, + const SheetSage2DecoderRuntimeOptions & options, + const SheetSage2DecoderWeights & weights, + int64_t batch, + int64_t memory_steps, + int64_t decoder_steps) + : backend_(execution.backend()), + backend_type_(execution.backend_type()), + threads_(std::max(1, execution.config().threads)), + config_(config), + options_(options), + weights_(weights), + batch_(batch), + memory_steps_(memory_steps), + decoder_steps_(decoder_steps) { + if (backend_ == nullptr || batch_ <= 0 || memory_steps_ <= 0 || decoder_steps_ <= 0) { + throw std::runtime_error("SheetSage2 decoder graph initialization failed"); + } + if (decoder_steps_ > config_.max_position_embeddings) { + throw std::runtime_error("SheetSage2 decoder steps exceed max position embeddings"); + } + build(); + } + + ~DecodeGraph() { + if (backend_ != nullptr && graph_ != nullptr) { + core::release_backend_graph_resources(backend_type_, backend_, graph_); + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + } + + bool matches(int64_t batch, int64_t memory_steps, int64_t decoder_steps) const noexcept { + return batch == batch_ && memory_steps == memory_steps_ && decoder_steps == decoder_steps_; + } + + std::vector run( + const std::vector & mixed_encoder_state, + const std::vector & decoder_input_ids) const { + if (static_cast(mixed_encoder_state.size()) != batch_ * memory_steps_ * config_.encoder_hidden_size) { + throw std::runtime_error("SheetSage2 mixed encoder state shape mismatch"); + } + if (static_cast(decoder_input_ids.size()) != batch_ * decoder_steps_) { + throw std::runtime_error("SheetSage2 decoder input id shape mismatch"); + } + core::write_tensor_f32(mixed_encoder_state_, mixed_encoder_state); + core::write_tensor_i32(decoder_input_ids_, decoder_input_ids); + core::write_tensor_i32(position_ids_, position_ids()); + core::set_backend_threads(backend_, threads_); + const ggml_status status = core::compute_backend_graph(backend_, graph_, nullptr, "models.sheetsage2.decoder"); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("SheetSage2 decoder graph compute failed"); + } + return core::read_tensor_f32(logits_); + } + +private: + std::vector position_ids() const { + std::vector ids(static_cast(batch_ * decoder_steps_)); + for (int64_t b = 0; b < batch_; ++b) { + for (int64_t t = 0; t < decoder_steps_; ++t) { + ids[static_cast(b * decoder_steps_ + t)] = static_cast(t + 2); + } + } + return ids; + } + + void build() { + ggml_init_params params{options_.graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("SheetSage2 decoder ggml context initialization failed"); + } + core::ModuleBuildContext input_ctx{ctx_.get(), "models.sheetsage2.decoder.inputs", backend_type_}; + mixed_encoder_state_ = core::make_tensor( + input_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({batch_, memory_steps_, config_.encoder_hidden_size})); + decoder_input_ids_ = core::make_tensor( + input_ctx, + GGML_TYPE_I32, + core::TensorShape::from_dims({batch_, decoder_steps_})); + position_ids_ = core::make_tensor( + input_ctx, + GGML_TYPE_I32, + core::TensorShape::from_dims({batch_, decoder_steps_})); + ggml_set_input(mixed_encoder_state_.tensor); + ggml_set_input(decoder_input_ids_.tensor); + ggml_set_input(position_ids_.tensor); + core::ModuleBuildContext build_ctx{ctx_.get(), "models.sheetsage2.decoder", backend_type_}; + auto logits = build_graph_output(build_ctx); + logits_ = logits.tensor; + ggml_set_output(logits_); + graph_ = ggml_new_graph_custom(ctx_.get(), 524288, false); + ggml_build_forward_expand(graph_, logits_); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (gallocr_ == nullptr || !ggml_gallocr_reserve(gallocr_, graph_) || !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + throw std::runtime_error("SheetSage2 decoder backend buffer allocation failed"); + } + } + + core::TensorValue build_graph_output(core::ModuleBuildContext & ctx) const { + auto memory = modules::LinearModule({config_.encoder_hidden_size, config_.hidden_size, true}).build( + ctx, + mixed_encoder_state_, + weights_.encoder_projection); + auto hidden = modules::EmbeddingModule({config_.vocab_size, config_.hidden_size}).build( + ctx, + decoder_input_ids_, + weights_.token_embedding); + auto positions = modules::EmbeddingModule({config_.max_position_embeddings + 2, config_.hidden_size}).build( + ctx, + position_ids_, + weights_.position_embedding); + hidden = modules::AddModule{}.build(ctx, hidden, positions); + hidden = modules::LayerNormModule({config_.hidden_size, config_.layer_norm_eps, true, true}).build( + ctx, + hidden, + weights_.layernorm_embedding); + for (const auto & layer : weights_.layers) { + hidden = decoder_layer(ctx, hidden, memory, layer, config_); + } + auto logits = modules::LinearModule({config_.hidden_size, config_.vocab_size, false}).build( + ctx, + hidden, + {weights_.token_embedding, std::nullopt}); + logits = modules::SliceModule({1, decoder_steps_ - 1, 1}).build(ctx, logits); + return core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, logits), + core::TensorShape::from_dims({config_.vocab_size})); + } + + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + SheetSage2DecoderConfig config_; + SheetSage2DecoderRuntimeOptions options_; + const SheetSage2DecoderWeights & weights_; + int64_t batch_ = 0; + int64_t memory_steps_ = 0; + int64_t decoder_steps_ = 0; + std::unique_ptr ctx_; + core::TensorValue mixed_encoder_state_; + core::TensorValue decoder_input_ids_; + core::TensorValue position_ids_; + ggml_tensor * logits_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; +}; + +class SheetSage2DecoderRuntime::Impl::CachedDecodeGraph { +public: + CachedDecodeGraph( + core::ExecutionContext & execution, + const SheetSage2DecoderConfig & config, + const SheetSage2DecoderRuntimeOptions & options, + const SheetSage2DecoderWeights & weights, + int64_t memory_steps, + int64_t cache_steps) + : backend_(execution.backend()), + backend_type_(execution.backend_type()), + threads_(std::max(1, execution.config().threads)), + config_(config), + options_(options), + weights_(weights), + memory_steps_(memory_steps), + cache_steps_(cache_steps), + head_dim_(config.hidden_size / config.num_attention_heads) { + if (backend_ == nullptr || memory_steps_ <= 0 || cache_steps_ <= 0) { + throw std::runtime_error("SheetSage2 cached decoder graph initialization failed"); + } + if (cache_steps_ > config_.max_position_embeddings) { + throw std::runtime_error("SheetSage2 cached decoder steps exceed max position embeddings"); + } + build(); + } + + ~CachedDecodeGraph() { + if (backend_ != nullptr && graph_ != nullptr) { + core::release_backend_graph_resources(backend_type_, backend_, graph_); + } + if (backend_ != nullptr && state_graph_ != nullptr) { + core::release_backend_graph_resources(backend_type_, backend_, state_graph_); + } + if (backend_ != nullptr && cross_graph_ != nullptr) { + core::release_backend_graph_resources(backend_type_, backend_, cross_graph_); + } + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + } + + bool matches(int64_t memory_steps, int64_t cache_steps) const noexcept { + return memory_steps == memory_steps_ && cache_steps == cache_steps_; + } + + std::vector prefill(const std::vector & tokens) { + const int64_t steps = static_cast(tokens.size()); + if (steps <= 0) { + throw std::runtime_error("SheetSage2 cached decoder prefill requires tokens"); + } + if (steps > cache_steps_) { + throw std::runtime_error("SheetSage2 cached decoder prefill exceeds cache capacity"); + } + for (int64_t i = 0; i + 1 < steps; ++i) { + run_step(tokens[static_cast(i)], false); + } + return run_step(tokens.back(), true); + } + + void reset(const std::vector & mixed_encoder_state) { + if (static_cast(mixed_encoder_state.size()) != memory_steps_ * config_.encoder_hidden_size) { + throw std::runtime_error("SheetSage2 cached decoder memory shape mismatch"); + } + core::write_tensor_f32(mixed_encoder_state_, mixed_encoder_state); + const auto cross_start = std::chrono::steady_clock::now(); + const ggml_status cross_status = + core::compute_backend_graph(backend_, cross_graph_, nullptr, "models.sheetsage2.decoder.cross_cache"); + if (cross_status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("SheetSage2 cached decoder cross-cache graph compute failed"); + } + engine::debug::timing_log_scalar( + "sheetsage2.decoder.cross_cache_ms", + engine::debug::elapsed_ms(cross_start)); + runtime::TransformerKVState empty; + empty.current_end = 0; + empty.layers.resize(static_cast(config_.decoder_layers)); + for (auto & layer : empty.layers) { + layer.valid_steps = 0; + } + step_cache_.import_state(empty); + attention_mask_values_.assign( + static_cast(config_.num_attention_heads * cache_steps_), + ggml_fp32_to_fp16(-INFINITY)); + ggml_backend_tensor_set( + attention_mask_.tensor, + attention_mask_values_.data(), + 0, + attention_mask_values_.size() * sizeof(ggml_fp16_t)); + } + + std::vector run_step(int32_t token, bool read_logits = true) { + if (step_cache_.valid_steps() >= cache_steps_) { + throw std::runtime_error("SheetSage2 cached decoder cache exhausted"); + } + core::write_tensor_i32(token_id_, std::vector{token}); + core::write_tensor_i32(position_id_, std::vector{static_cast(step_cache_.current_end() + 2)}); + core::write_tensor_i32(cache_slot_, std::vector{static_cast(step_cache_.valid_steps())}); + const size_t valid = static_cast(step_cache_.valid_steps()); + for (int64_t head = 0; head < config_.num_attention_heads; ++head) { + attention_mask_values_[static_cast(head * cache_steps_) + valid] = ggml_fp32_to_fp16(0.0F); + } + ggml_backend_tensor_set( + attention_mask_.tensor, + attention_mask_values_.data(), + 0, + attention_mask_values_.size() * sizeof(ggml_fp16_t)); + core::set_backend_threads(backend_, threads_); + const ggml_status status = core::compute_backend_graph( + backend_, + read_logits ? graph_ : state_graph_, + nullptr, + read_logits ? "models.sheetsage2.decoder.cached_step" : "models.sheetsage2.decoder.cached_state_step"); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("SheetSage2 cached decoder graph compute failed"); + } + ggml_backend_synchronize(backend_); + step_cache_.advance_after_direct_append(1); + return read_logits ? core::read_tensor_f32(logits_) : std::vector{}; + } + +private: + void build() { + ggml_init_params params{options_.graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("SheetSage2 cached decoder ggml context initialization failed"); + } + core::ModuleBuildContext input_ctx{ctx_.get(), "models.sheetsage2.decoder.cached.inputs", backend_type_}; + mixed_encoder_state_ = core::make_tensor( + input_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, memory_steps_, config_.encoder_hidden_size})); + token_id_ = core::make_tensor( + input_ctx, + GGML_TYPE_I32, + core::TensorShape::from_dims({1, 1})); + position_id_ = core::make_tensor( + input_ctx, + GGML_TYPE_I32, + core::TensorShape::from_dims({1, 1})); + cache_slot_ = core::make_tensor( + input_ctx, + GGML_TYPE_I32, + core::TensorShape::from_dims({1})); + attention_mask_ = core::make_tensor( + input_ctx, + GGML_TYPE_F16, + core::TensorShape::from_dims({1, config_.num_attention_heads, 1, cache_steps_})); + ggml_set_input(mixed_encoder_state_.tensor); + ggml_set_input(token_id_.tensor); + ggml_set_input(position_id_.tensor); + ggml_set_input(cache_slot_.tensor); + ggml_set_input(attention_mask_.tensor); + + std::vector cache_keys; + std::vector cache_values; + cache_keys.reserve(static_cast(config_.decoder_layers)); + cache_values.reserve(static_cast(config_.decoder_layers)); + cross_keys_.reserve(static_cast(config_.decoder_layers)); + cross_values_.reserve(static_cast(config_.decoder_layers)); + for (int64_t layer = 0; layer < config_.decoder_layers; ++layer) { + cache_keys.push_back(core::make_tensor( + input_ctx, + GGML_TYPE_F16, + core::TensorShape::from_dims({1, cache_steps_, config_.num_attention_heads, head_dim_}))); + cache_values.push_back(core::make_tensor( + input_ctx, + GGML_TYPE_F16, + core::TensorShape::from_dims({1, cache_steps_, config_.num_attention_heads, head_dim_}))); + ggml_set_input(cache_keys.back().tensor); + ggml_set_input(cache_values.back().tensor); + cross_keys_.push_back(core::make_tensor( + input_ctx, + GGML_TYPE_F16, + core::TensorShape::from_dims({1, config_.num_attention_heads, memory_steps_, head_dim_}))); + cross_values_.push_back(core::make_tensor( + input_ctx, + GGML_TYPE_F16, + core::TensorShape::from_dims({1, config_.num_attention_heads, memory_steps_, head_dim_}))); + } + runtime::TransformerKVCacheOptions cache_options; + cache_options.allow_f16_storage = true; + step_cache_ = runtime::TransformerKVCache( + cache_steps_, + config_.num_attention_heads * head_dim_, + std::move(cache_keys), + std::move(cache_values), + cache_options); + + core::ModuleBuildContext build_ctx{ctx_.get(), "models.sheetsage2.decoder.cached", backend_type_}; + build_cross_cache_graph(build_ctx); + auto hidden = build_hidden_output(build_ctx); + hidden_ = hidden.tensor; + ggml_set_output(hidden_); + state_graph_ = ggml_new_graph_custom(ctx_.get(), 524288, false); + ggml_build_forward_expand(state_graph_, hidden_); + auto logits = build_logits_output(build_ctx, hidden); + logits_ = logits.tensor; + ggml_set_output(logits_); + graph_ = ggml_new_graph_custom(ctx_.get(), 524288, false); + ggml_build_forward_expand(graph_, logits_); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), backend_); + if (buffer_ == nullptr) { + throw std::runtime_error("SheetSage2 cached decoder backend buffer allocation failed"); + } + } + + void build_cross_cache_graph(core::ModuleBuildContext & ctx) { + auto memory = modules::LinearModule({config_.encoder_hidden_size, config_.hidden_size, true}).build( + ctx, + mixed_encoder_state_, + weights_.encoder_projection); + cross_graph_ = ggml_new_graph_custom(ctx.ggml, 131072, false); + for (int64_t layer = 0; layer < config_.decoder_layers; ++layer) { + const auto kv = cross_attention_key_value( + ctx, + memory, + weights_.layers[static_cast(layer)].encoder_attn, + config_.hidden_size, + config_.num_attention_heads); + auto key_copy = core::wrap_tensor( + ggml_cpy(ctx.ggml, kv.key.tensor, cross_keys_[static_cast(layer)].tensor), + cross_keys_[static_cast(layer)].shape, + GGML_TYPE_F16); + auto value_copy = core::wrap_tensor( + ggml_cpy(ctx.ggml, kv.value.tensor, cross_values_[static_cast(layer)].tensor), + cross_values_[static_cast(layer)].shape, + GGML_TYPE_F16); + ggml_set_output(key_copy.tensor); + ggml_set_output(value_copy.tensor); + ggml_build_forward_expand(cross_graph_, key_copy.tensor); + ggml_build_forward_expand(cross_graph_, value_copy.tensor); + } + } + + core::TensorValue build_hidden_output(core::ModuleBuildContext & ctx) { + auto hidden = modules::EmbeddingModule({config_.vocab_size, config_.hidden_size}).build( + ctx, + token_id_, + weights_.token_embedding); + auto positions = modules::EmbeddingModule({config_.max_position_embeddings + 2, config_.hidden_size}).build( + ctx, + position_id_, + weights_.position_embedding); + hidden = modules::AddModule{}.build(ctx, hidden, positions); + hidden = modules::LayerNormModule({config_.hidden_size, config_.layer_norm_eps, true, true}).build( + ctx, + hidden, + weights_.layernorm_embedding); + for (int64_t layer = 0; layer < config_.decoder_layers; ++layer) { + auto out = cached_decoder_layer_step( + ctx, + hidden, + {cross_keys_[static_cast(layer)], cross_values_[static_cast(layer)]}, + step_cache_.key_tensor(static_cast(layer)), + step_cache_.value_tensor(static_cast(layer)), + cache_slot_, + attention_mask_, + weights_.layers[static_cast(layer)], + config_); + hidden = out.hidden; + } + return core::ensure_backend_addressable_layout(ctx, hidden); + } + + core::TensorValue build_logits_output(core::ModuleBuildContext & ctx, const core::TensorValue & hidden) { + auto logits = modules::LinearModule({config_.hidden_size, config_.vocab_size, false}).build( + ctx, + hidden, + {weights_.token_embedding, std::nullopt}); + return core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, logits), + core::TensorShape::from_dims({config_.vocab_size})); + } + + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + SheetSage2DecoderConfig config_; + SheetSage2DecoderRuntimeOptions options_; + const SheetSage2DecoderWeights & weights_; + int64_t memory_steps_ = 0; + int64_t cache_steps_ = 0; + int64_t head_dim_ = 0; + std::unique_ptr ctx_; + core::TensorValue mixed_encoder_state_; + core::TensorValue token_id_; + core::TensorValue position_id_; + core::TensorValue cache_slot_; + core::TensorValue attention_mask_; + std::vector cross_keys_; + std::vector cross_values_; + std::vector attention_mask_values_; + runtime::TransformerKVCache step_cache_; + ggml_tensor * hidden_ = nullptr; + ggml_tensor * logits_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_cgraph * state_graph_ = nullptr; + ggml_cgraph * cross_graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; +}; + +Mert2EncoderRuntime::Mert2EncoderRuntime( + std::shared_ptr source, + core::ExecutionContext & execution, + SheetSage2DecoderConfig config, + SheetSage2DecoderRuntimeOptions options) + : impl_(std::make_unique(std::move(source), execution, config, options)) {} + +Mert2EncoderRuntime::~Mert2EncoderRuntime() = default; +Mert2EncoderRuntime::Mert2EncoderRuntime(Mert2EncoderRuntime &&) noexcept = default; +Mert2EncoderRuntime & Mert2EncoderRuntime::operator=(Mert2EncoderRuntime &&) noexcept = default; + +void Mert2EncoderRuntime::prepare(int64_t batch, int64_t mel_frames) { + const auto & weights = impl_->require_encoder_weights(); + if (!impl_->encoder_graph || !impl_->encoder_graph->matches(batch, mel_frames)) { + impl_->encoder_graph = std::make_unique( + *impl_->execution, + impl_->config, + impl_->options, + weights, + batch, + mel_frames); + } +} + +std::vector Mert2EncoderRuntime::encode_mel( + const std::vector & normalized_mel, + int64_t mel_frames) { + if (mel_frames <= 0) { + throw std::runtime_error("MERT2 mel frames must be positive"); + } + if (static_cast(normalized_mel.size()) % (mel_frames * impl_->config.mel_bins) != 0) { + throw std::runtime_error("MERT2 normalized mel does not divide into batches"); + } + const int64_t batch = static_cast(normalized_mel.size()) / (mel_frames * impl_->config.mel_bins); + prepare(batch, mel_frames); + return impl_->encoder_graph->run(normalized_mel); +} + +void Mert2EncoderRuntime::release_runtime_graphs() { + impl_->encoder_graph.reset(); +} + +SheetSage2DecoderRuntime::SheetSage2DecoderRuntime( + std::shared_ptr source, + core::ExecutionContext & execution, + SheetSage2DecoderConfig config, + SheetSage2DecoderRuntimeOptions options) + : impl_(std::make_unique(std::move(source), execution, config, options)) {} + +SheetSage2DecoderRuntime::~SheetSage2DecoderRuntime() = default; +SheetSage2DecoderRuntime::SheetSage2DecoderRuntime(SheetSage2DecoderRuntime &&) noexcept = default; +SheetSage2DecoderRuntime & SheetSage2DecoderRuntime::operator=(SheetSage2DecoderRuntime &&) noexcept = default; + +void SheetSage2DecoderRuntime::prepare(int64_t batch, int64_t memory_steps, int64_t decoder_steps) { + const auto & weights = impl_->require_weights(); + if (!impl_->graph || !impl_->graph->matches(batch, memory_steps, decoder_steps)) { + impl_->graph = std::make_unique( + *impl_->execution, + impl_->config, + impl_->options, + weights, + batch, + memory_steps, + decoder_steps); + } +} + +std::vector SheetSage2DecoderRuntime::decode_logits( + const std::vector & mixed_encoder_state, + int64_t memory_steps, + const std::vector & decoder_input_ids) { + if (memory_steps <= 0) { + throw std::runtime_error("SheetSage2 memory steps must be positive"); + } + if (decoder_input_ids.empty()) { + throw std::runtime_error("SheetSage2 decoder input ids must not be empty"); + } + const int64_t memory_values_per_batch = memory_steps * impl_->config.encoder_hidden_size; + if (static_cast(mixed_encoder_state.size()) % memory_values_per_batch != 0) { + throw std::runtime_error("SheetSage2 mixed encoder state does not divide into batches"); + } + const int64_t batch = static_cast(mixed_encoder_state.size()) / memory_values_per_batch; + if (static_cast(decoder_input_ids.size()) % batch != 0) { + throw std::runtime_error("SheetSage2 decoder input ids do not divide by batch"); + } + const int64_t decoder_steps = static_cast(decoder_input_ids.size()) / batch; + prepare(batch, memory_steps, decoder_steps); + return impl_->graph->run(mixed_encoder_state, decoder_input_ids); +} + +void SheetSage2DecoderRuntime::reset_cached_decode( + const std::vector & mixed_encoder_state, + int64_t memory_steps, + int64_t cache_steps) { + if (memory_steps <= 0 || cache_steps <= 0) { + throw std::runtime_error("SheetSage2 cached decode requires positive memory/cache steps"); + } + const int64_t expected = memory_steps * impl_->config.encoder_hidden_size; + if (static_cast(mixed_encoder_state.size()) != expected) { + throw std::runtime_error("SheetSage2 cached decode memory shape mismatch"); + } + const auto & weights = impl_->require_weights(); + if (!impl_->cached_graph || !impl_->cached_graph->matches(memory_steps, cache_steps)) { + impl_->cached_graph = std::make_unique( + *impl_->execution, + impl_->config, + impl_->options, + weights, + memory_steps, + cache_steps); + } + impl_->cached_graph->reset(mixed_encoder_state); +} + +std::vector SheetSage2DecoderRuntime::decode_cached_step(int32_t token) { + if (!impl_->cached_graph) { + throw std::runtime_error("SheetSage2 cached decode graph is not prepared"); + } + return impl_->cached_graph->run_step(token); +} + +std::vector SheetSage2DecoderRuntime::prefill_cached_decode(const std::vector & token_ids) { + if (!impl_->cached_graph) { + throw std::runtime_error("SheetSage2 cached decode graph is not prepared"); + } + return impl_->cached_graph->prefill(token_ids); +} + +void SheetSage2DecoderRuntime::release_runtime_graphs() { + impl_->graph.reset(); + impl_->cached_graph.reset(); +} + +} // namespace engine::models::sheetsage diff --git a/src/models/sheetsage/session.cpp b/src/models/sheetsage/session.cpp new file mode 100644 index 00000000..600046fb --- /dev/null +++ b/src/models/sheetsage/session.cpp @@ -0,0 +1,1114 @@ +#include "engine/models/sheetsage/session.h" +#include "engine/models/sheetsage/processing.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/audio/dsp.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/model_spec/package.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/runtime/spec_backed_model.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::sheetsage { +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr const char * kFamily = "sheetsage2"; +constexpr int64_t kTimeHz = 100; +constexpr int64_t kPromptCapacity = 256; +constexpr int64_t kMaxSubbeatShift = 256; +constexpr double kDefaultOverlapSeconds = 200.0; +constexpr double kDefaultLookaheadSeconds = 100.0; +constexpr double kDefaultStepSeconds = 0.125; + +constexpr std::array kPromptNames = { + "timestamp", + "downbeat_meter", + "structure", + "key", + "chord_majmin", + "chord_full", + "melody_vocal", + "melody_full", +}; + +struct SheetSage2Tokenizer { + int64_t audio_seconds = 300; + int64_t sos = 1; + int64_t eos = 2; + int64_t out = 3; + int64_t prompt_start = 4; + int64_t prompt_end = prompt_start + kPromptCapacity; + int64_t subbeat_start = prompt_end; + int64_t subbeat_end = subbeat_start + kMaxSubbeatShift + 1; + int64_t time_start = subbeat_end; + int64_t time_end = time_start + audio_seconds * kTimeHz; + int64_t meter_start = time_end; + int64_t meter_end = meter_start + 32 * 6; + int64_t eighth_start = meter_end; + int64_t eighth_end = eighth_start + 256; + int64_t structure_start = eighth_end; + int64_t structure_end = structure_start + sheetsage2_structure_label_count(); + int64_t key_start = structure_end; + int64_t key_end = key_start + 24; + int64_t majmin_chord_start = key_end; + int64_t majmin_chord_end = majmin_chord_start + 25; + int64_t full_chord_start = majmin_chord_end; + int64_t full_chord_end = full_chord_start + 361; + int64_t pitch_start = full_chord_end; + int64_t pitch_end = pitch_start + 256; + int64_t duration_start = pitch_end; + int64_t duration_end = duration_start + sheetsage2_duration_bin_count(); + int64_t vocab_size = duration_end; + + explicit SheetSage2Tokenizer(int64_t seconds) : audio_seconds(std::max(1, seconds)) { + time_end = time_start + audio_seconds * kTimeHz; + meter_start = time_end; + meter_end = meter_start + 32 * 6; + eighth_start = meter_end; + eighth_end = eighth_start + 256; + structure_start = eighth_end; + structure_end = structure_start + sheetsage2_structure_label_count(); + key_start = structure_end; + key_end = key_start + 24; + majmin_chord_start = key_end; + majmin_chord_end = majmin_chord_start + 25; + full_chord_start = majmin_chord_end; + full_chord_end = full_chord_start + 361; + pitch_start = full_chord_end; + pitch_end = pitch_start + 256; + duration_start = pitch_end; + duration_end = duration_start + sheetsage2_duration_bin_count(); + vocab_size = duration_end; + } +}; + +enum class SheetSage2TokenType { + Special, + Subbeat, + Time, + Meter, + Eighth, + Structure, + Key, + Chord, + Pitch, + Duration, +}; + +struct SheetSage2GrammarState { + int payload_count = 0; + bool in_shift = true; + int shift_run = 0; + int last_field_index = -1; + enum class Incomplete { + None, + RhythmAfterMeter, + MelodyAfterPitch, + } incomplete = Incomplete::None; +}; + +struct SheetSage2Window { + int64_t index = 0; + int64_t start_sample = 0; + int64_t end_sample = 0; + double start = 0.0; + double end = 0.0; + double accept_start = 0.0; + double accept_end = 0.0; + double prefix_end = 0.0; + std::optional generation_stop; +}; + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("SheetSage2 session requires assets"); + } + return assets; +} + +std::shared_ptr require_contract( + std::shared_ptr contract) { + if (contract == nullptr) { + throw std::runtime_error("SheetSage2 session requires a model contract"); + } + return contract; +} + +SheetSage2DecoderConfig parse_decoder_config(const assets::ResourceBundle & resources) { + const auto root = resources.parse_json("config"); + SheetSage2DecoderConfig config; + config.vocab_size = io::json::optional_i64(root, "vocab_size", config.vocab_size); + config.hidden_size = io::json::optional_i64(root, "hidden_size", config.hidden_size); + config.intermediate_size = io::json::optional_i64(root, "intermediate_size", config.intermediate_size); + config.decoder_layers = io::json::optional_i64(root, "decoder_layers", config.decoder_layers); + config.num_attention_heads = io::json::optional_i64(root, "num_attention_heads", config.num_attention_heads); + config.max_position_embeddings = io::json::optional_i64(root, "max_output_seq_len", config.max_position_embeddings); + config.pad_token_id = io::json::optional_i64(root, "pad_token_id", config.pad_token_id); + config.layer_norm_eps = io::json::optional_f32(root, "layer_norm_eps", config.layer_norm_eps); + if (const auto * backbone = root.find("backbone_config")) { + config.encoder_hidden_size = io::json::optional_i64(*backbone, "hidden_size", config.encoder_hidden_size); + config.encoder_intermediate_size = io::json::optional_i64(*backbone, "intermediate_size", config.encoder_intermediate_size); + config.encoder_layers = io::json::optional_i64(*backbone, "num_hidden_layers", config.encoder_layers); + config.encoder_attention_heads = io::json::optional_i64(*backbone, "num_attention_heads", config.encoder_attention_heads); + config.sampling_rate = io::json::optional_i64(*backbone, "sampling_rate", config.sampling_rate); + config.n_fft = io::json::optional_i64(*backbone, "n_fft", config.n_fft); + config.win_length = io::json::optional_i64(*backbone, "win_length", config.win_length); + config.hop_length = io::json::optional_i64(*backbone, "hop_length", config.hop_length); + config.mel_bins = io::json::optional_i64(*backbone, "num_mel_bins", config.mel_bins); + config.input_audio_length_samples = + io::json::optional_i64(*backbone, "input_audio_length", config.input_audio_length_samples); + config.encoder_layer_norm_eps = io::json::optional_f32(*backbone, "layer_norm_eps", config.encoder_layer_norm_eps); + config.subsampling_layer_norm_eps = + io::json::optional_f32(*backbone, "subsampling_layer_norm_eps", config.subsampling_layer_norm_eps); + config.rotary_embedding_base = + io::json::optional_f32(*backbone, "rotary_embedding_base", config.rotary_embedding_base); + config.conformer_conv_kernel_size = + io::json::optional_i64(*backbone, "conv_depthwise_kernel_size", config.conformer_conv_kernel_size); + } + return config; +} + +SheetSage2DecoderRuntimeOptions decoder_options_from_session_options(const runtime::SessionOptions & options) { + SheetSage2DecoderRuntimeOptions out; + out.weight_storage_type = runtime::parse_tensor_storage_option( + options.options, + "sheetsage2.weight_type", + assets::TensorStorageType::Native, + { + assets::TensorStorageType::Native, + assets::TensorStorageType::F32, + assets::TensorStorageType::F16, + assets::TensorStorageType::BF16, + assets::TensorStorageType::Q4_0, + assets::TensorStorageType::Q4_K, + }); + out.weight_context_bytes = + runtime::parse_size_mb_option(options.options, {"sheetsage2.weight_context_mb"}, out.weight_context_bytes); + out.graph_arena_bytes = + runtime::parse_size_mb_option(options.options, {"sheetsage2.decoder_graph_arena_mb"}, out.graph_arena_bytes); + return out; +} + +std::vector prompt_prefix(const SheetSage2Tokenizer & tokenizer) { + return { + static_cast(tokenizer.sos), + static_cast(tokenizer.prompt_start), + static_cast(tokenizer.prompt_start + 1), + static_cast(tokenizer.prompt_start + 2), + static_cast(tokenizer.prompt_start + 3), + static_cast(tokenizer.prompt_start + 5), + static_cast(tokenizer.prompt_start + 7), + static_cast(tokenizer.out), + }; +} + +SheetSage2TokenType token_type(const SheetSage2Tokenizer & tokenizer, int64_t token) { + if (token == tokenizer.sos || token == tokenizer.eos || token == tokenizer.out || + (token >= tokenizer.prompt_start && token < tokenizer.prompt_start + static_cast(kPromptNames.size()))) { + return SheetSage2TokenType::Special; + } + if (token >= tokenizer.subbeat_start && token < tokenizer.subbeat_end) { + return SheetSage2TokenType::Subbeat; + } + if (token >= tokenizer.time_start && token < tokenizer.time_end) { + return SheetSage2TokenType::Time; + } + if (token >= tokenizer.meter_start && token < tokenizer.meter_end) { + return SheetSage2TokenType::Meter; + } + if (token >= tokenizer.eighth_start && token < tokenizer.eighth_end) { + return SheetSage2TokenType::Eighth; + } + if (token >= tokenizer.structure_start && token < tokenizer.structure_end) { + return SheetSage2TokenType::Structure; + } + if (token >= tokenizer.key_start && token < tokenizer.key_end) { + return SheetSage2TokenType::Key; + } + if ((token >= tokenizer.majmin_chord_start && token < tokenizer.majmin_chord_end) || + (token >= tokenizer.full_chord_start && token < tokenizer.full_chord_end)) { + return SheetSage2TokenType::Chord; + } + if (token >= tokenizer.pitch_start && token < tokenizer.pitch_end) { + return SheetSage2TokenType::Pitch; + } + if (token >= tokenizer.duration_start && token < tokenizer.duration_end) { + return SheetSage2TokenType::Duration; + } + return SheetSage2TokenType::Special; +} + +bool update_grammar(const SheetSage2Tokenizer & tokenizer, SheetSage2GrammarState & state, int64_t token) { + if (token == tokenizer.eos) { + return true; + } + const auto type = token_type(tokenizer, token); + if (type == SheetSage2TokenType::Subbeat) { + if (!state.in_shift && state.payload_count > 0) { + state.payload_count = 0; + state.last_field_index = -1; + state.incomplete = SheetSage2GrammarState::Incomplete::None; + } + state.in_shift = true; + ++state.shift_run; + return false; + } + state.in_shift = false; + state.shift_run = 0; + ++state.payload_count; + switch (type) { + case SheetSage2TokenType::Time: + state.last_field_index = 0; + state.incomplete = SheetSage2GrammarState::Incomplete::None; + break; + case SheetSage2TokenType::Meter: + state.last_field_index = 1; + state.incomplete = SheetSage2GrammarState::Incomplete::RhythmAfterMeter; + break; + case SheetSage2TokenType::Eighth: + state.last_field_index = 1; + state.incomplete = SheetSage2GrammarState::Incomplete::None; + break; + case SheetSage2TokenType::Structure: + state.last_field_index = 2; + state.incomplete = SheetSage2GrammarState::Incomplete::None; + break; + case SheetSage2TokenType::Key: + state.last_field_index = 3; + state.incomplete = SheetSage2GrammarState::Incomplete::None; + break; + case SheetSage2TokenType::Chord: + state.last_field_index = 4; + state.incomplete = SheetSage2GrammarState::Incomplete::None; + break; + case SheetSage2TokenType::Pitch: + state.last_field_index = 5; + state.incomplete = SheetSage2GrammarState::Incomplete::MelodyAfterPitch; + break; + case SheetSage2TokenType::Duration: + state.last_field_index = 5; + state.incomplete = SheetSage2GrammarState::Incomplete::None; + break; + case SheetSage2TokenType::Special: + case SheetSage2TokenType::Subbeat: + throw std::runtime_error("SheetSage2 generated invalid grammar token"); + } + return false; +} + +int32_t select_next_token( + const std::vector & logits, + const SheetSage2Tokenizer & tokenizer, + const SheetSage2GrammarState & state) { + int32_t best = static_cast(tokenizer.eos); + float best_value = -std::numeric_limits::infinity(); + const int64_t start = static_cast(logits.size()) - tokenizer.vocab_size; + if (start < 0) { + throw std::runtime_error("SheetSage2 decoder logits are smaller than vocabulary"); + } + const auto scan_token = [&](int64_t token) { + const float value = logits[static_cast(start + token)]; + if (value > best_value) { + best_value = value; + best = static_cast(token); + } + }; + const auto scan_range = [&](int64_t begin, int64_t end) { + for (int64_t token = begin; token < end; ++token) { + scan_token(token); + } + }; + if (state.payload_count > 0) { + scan_token(tokenizer.eos); + } + if ((state.payload_count > 0 || state.in_shift) && state.shift_run < 4) { + scan_range(tokenizer.subbeat_start, tokenizer.subbeat_end); + } + if (state.incomplete == SheetSage2GrammarState::Incomplete::RhythmAfterMeter) { + scan_range(tokenizer.eighth_start, tokenizer.eighth_end); + return best; + } + if (state.incomplete == SheetSage2GrammarState::Incomplete::MelodyAfterPitch) { + scan_range(tokenizer.pitch_start, tokenizer.pitch_end); + scan_range(tokenizer.duration_start, tokenizer.duration_end); + return best; + } + if (state.last_field_index < 0) { + scan_range(tokenizer.time_start, tokenizer.time_end); + } + if (state.last_field_index < 1) { + scan_range(tokenizer.meter_start, tokenizer.meter_end); + scan_range(tokenizer.eighth_start, tokenizer.eighth_end); + } + if (state.last_field_index < 2) { + scan_range(tokenizer.structure_start, tokenizer.structure_end); + } + if (state.last_field_index < 3) { + scan_range(tokenizer.key_start, tokenizer.key_end); + } + if (state.last_field_index < 4) { + scan_range(tokenizer.full_chord_start, tokenizer.full_chord_end); + } + if (state.last_field_index <= 5) { + scan_range(tokenizer.pitch_start, tokenizer.pitch_end); + } + return best; +} + +std::vector decode_events( + const SheetSage2Tokenizer & tokenizer, + const std::vector & tokens) { + std::vector events; + auto out_it = std::find(tokens.begin(), tokens.end(), static_cast(tokenizer.out)); + if (out_it == tokens.end()) { + return events; + } + int64_t current_step = 0; + size_t pos = static_cast(std::distance(tokens.begin(), out_it)) + 1; + while (pos < tokens.size()) { + int64_t token = tokens[pos]; + if (token == tokenizer.eos) { + break; + } + if (token_type(tokenizer, token) != SheetSage2TokenType::Subbeat) { + ++pos; + continue; + } + int64_t shift = 0; + while (pos < tokens.size() && token_type(tokenizer, tokens[pos]) == SheetSage2TokenType::Subbeat) { + shift += tokens[pos] - tokenizer.subbeat_start; + ++pos; + } + current_step += shift; + SheetSage2Event event; + event.subbeat = current_step; + while (pos < tokens.size()) { + token = tokens[pos]; + const auto type = token_type(tokenizer, token); + if (token == tokenizer.eos || type == SheetSage2TokenType::Subbeat) { + break; + } + if (type == SheetSage2TokenType::Time) { + event.timestamp = static_cast(token - tokenizer.time_start) / static_cast(kTimeHz); + event.timestamp_tokens.push_back(static_cast(token)); + } else if (type == SheetSage2TokenType::Meter) { + const int64_t index = token - tokenizer.meter_start; + static constexpr std::array denominators = {1, 2, 4, 8, 16, 32}; + event.meter = {static_cast(index / 6 + 1), denominators[static_cast(index % 6)]}; + event.rhythm_tokens.push_back(static_cast(token)); + } else if (type == SheetSage2TokenType::Eighth) { + event.eighth_position = token - tokenizer.eighth_start; + event.rhythm_tokens.push_back(static_cast(token)); + } else if (type == SheetSage2TokenType::Structure) { + const int64_t index = token - tokenizer.structure_start; + event.structure = sheetsage2_structure_label(index); + event.structure_tokens.push_back(static_cast(token)); + } else if (type == SheetSage2TokenType::Key) { + const int64_t index = token - tokenizer.key_start; + event.key = sheetsage2_key_label(index); + event.key_tokens.push_back(static_cast(token)); + } else if (type == SheetSage2TokenType::Chord) { + if (token >= tokenizer.full_chord_start && token < tokenizer.full_chord_end) { + const int64_t index = token - tokenizer.full_chord_start; + event.chord = sheetsage2_chord_label(true, index); + } else { + const int64_t index = token - tokenizer.majmin_chord_start; + event.chord = sheetsage2_chord_label(false, index); + } + event.chord_tokens.push_back(static_cast(token)); + } else if (type == SheetSage2TokenType::Pitch) { + const int pitch_id = static_cast(token - tokenizer.pitch_start); + int duration_bin = 0; + event.melody_tokens.push_back(static_cast(token)); + if (pos + 1 < tokens.size() && token_type(tokenizer, tokens[pos + 1]) == SheetSage2TokenType::Duration) { + const int64_t bin = tokens[pos + 1] - tokenizer.duration_start; + duration_bin = static_cast(std::clamp( + bin, + 0, + sheetsage2_duration_bin_count() - 1)); + event.melody_tokens.push_back(tokens[pos + 1]); + ++pos; + } + event.notes.push_back(sheetsage2_note_from_pitch_duration(pitch_id, duration_bin)); + } + ++pos; + } + events.push_back(std::move(event)); + } + return events; +} + +std::vector subbeat_shift_tokens(const SheetSage2Tokenizer & tokenizer, int64_t shift) { + if (shift < 0) { + throw std::runtime_error("SheetSage2 prefix events are not sorted by subbeat"); + } + std::vector tokens; + while (shift > kMaxSubbeatShift) { + tokens.push_back(static_cast(tokenizer.subbeat_start + kMaxSubbeatShift)); + shift -= kMaxSubbeatShift; + } + tokens.push_back(static_cast(tokenizer.subbeat_start + shift)); + return tokens; +} + +std::vector encode_events( + const SheetSage2Tokenizer & tokenizer, + const std::vector & events, + bool include_eos) { + auto ids = prompt_prefix(tokenizer); + int64_t previous_step = 0; + for (const auto & event : events) { + auto shift = subbeat_shift_tokens(tokenizer, event.subbeat - previous_step); + ids.insert(ids.end(), shift.begin(), shift.end()); + previous_step = event.subbeat; + ids.insert(ids.end(), event.timestamp_tokens.begin(), event.timestamp_tokens.end()); + ids.insert(ids.end(), event.rhythm_tokens.begin(), event.rhythm_tokens.end()); + ids.insert(ids.end(), event.structure_tokens.begin(), event.structure_tokens.end()); + ids.insert(ids.end(), event.key_tokens.begin(), event.key_tokens.end()); + ids.insert(ids.end(), event.chord_tokens.begin(), event.chord_tokens.end()); + ids.insert(ids.end(), event.melody_tokens.begin(), event.melody_tokens.end()); + } + if (include_eos) { + ids.push_back(static_cast(tokenizer.eos)); + } + return ids; +} + +float event_time_value(const SheetSage2Event & event) { + return event.time; +} + +std::function make_event_time_lookup( + const std::vector & events, + double target_seconds) { + std::vector> anchors; + for (const auto & event : events) { + if (event.timestamp.has_value()) { + anchors.push_back({event.subbeat, *event.timestamp}); + } + } + if (anchors.empty()) { + return [target_seconds](int64_t step) { + return std::min(target_seconds, std::max(0.0, static_cast(step) * kDefaultStepSeconds)); + }; + } + std::sort(anchors.begin(), anchors.end()); + anchors.erase( + std::unique( + anchors.begin(), + anchors.end(), + [](const auto & a, const auto & b) { return a.first == b.first; }), + anchors.end()); + double step_seconds = kDefaultStepSeconds; + if (anchors.size() >= 2) { + std::vector slopes; + slopes.reserve(anchors.size() - 1); + for (size_t i = 1; i < anchors.size(); ++i) { + const double step_delta = static_cast(std::max(1, anchors[i].first - anchors[i - 1].first)); + slopes.push_back((anchors[i].second - anchors[i - 1].second) / step_delta); + } + std::sort(slopes.begin(), slopes.end()); + step_seconds = slopes[slopes.size() / 2]; + if (!std::isfinite(step_seconds) || step_seconds <= 0.0) { + step_seconds = kDefaultStepSeconds; + } + } + return [anchors, target_seconds, step_seconds](int64_t step) { + const double value = static_cast(step); + if (value <= static_cast(anchors.front().first)) { + return std::clamp( + anchors.front().second + (value - static_cast(anchors.front().first)) * step_seconds, + 0.0, + target_seconds); + } + if (value >= static_cast(anchors.back().first)) { + return std::clamp( + anchors.back().second + (value - static_cast(anchors.back().first)) * step_seconds, + 0.0, + target_seconds); + } + for (size_t i = 1; i < anchors.size(); ++i) { + if (value <= static_cast(anchors[i].first)) { + const double left_step = static_cast(anchors[i - 1].first); + const double right_step = static_cast(anchors[i].first); + const double alpha = (value - left_step) / std::max(1.0, right_step - left_step); + return anchors[i - 1].second + alpha * (anchors[i].second - anchors[i - 1].second); + } + } + return anchors.back().second; + }; +} + +std::vector sliding_window_plan( + int64_t audio_samples, + int64_t sample_rate, + int64_t window_samples) { + if (audio_samples <= 0 || sample_rate <= 0 || window_samples <= 0) { + throw std::runtime_error("SheetSage2 sliding window requires positive audio shape"); + } + const double duration = static_cast(audio_samples) / static_cast(sample_rate); + const double window_seconds = static_cast(window_samples) / static_cast(sample_rate); + const double overlap_seconds = std::min(kDefaultOverlapSeconds, std::max(0.0, window_seconds - 1.0)); + const double lookahead_seconds = std::min(kDefaultLookaheadSeconds, overlap_seconds); + if (!(0.0 <= lookahead_seconds && lookahead_seconds <= overlap_seconds && overlap_seconds < window_seconds)) { + throw std::runtime_error("SheetSage2 invalid sliding window overlap"); + } + const double hop = window_seconds - overlap_seconds; + double start = 0.0; + double accepted = 0.0; + int64_t index = 0; + std::vector out; + while (true) { + const bool last = start + window_seconds >= duration - 1.0e-6; + const double accept_end = last ? duration : start + window_seconds - lookahead_seconds; + SheetSage2Window window; + window.index = index++; + window.start = start; + window.end = std::min(duration, start + window_seconds); + window.accept_start = accepted; + window.accept_end = accept_end; + window.prefix_end = accepted; + window.generation_stop = last ? std::nullopt : std::optional(window_seconds - lookahead_seconds); + window.start_sample = static_cast(std::llround(start * static_cast(sample_rate))); + window.end_sample = std::min( + audio_samples, + window.start_sample + window_samples); + out.push_back(window); + if (last) { + return out; + } + accepted = accept_end; + start = std::min(start + hop, duration - window_seconds); + } +} + +std::vector slice_window_audio( + const std::vector & audio, + int64_t start_sample, + int64_t window_samples) { + std::vector out(static_cast(window_samples), 0.0F); + if (start_sample < 0 || start_sample > static_cast(audio.size())) { + throw std::runtime_error("SheetSage2 sliding window start is outside audio"); + } + const int64_t count = std::min(window_samples, static_cast(audio.size()) - start_sample); + if (count > 0) { + std::copy_n( + audio.begin() + static_cast(start_sample), + static_cast(count), + out.begin()); + } + return out; +} + +void set_event_local_timestamp( + SheetSage2Event & event, + const SheetSage2Tokenizer & tokenizer, + double local_time) { + int64_t time_id = static_cast(std::llround(local_time * static_cast(kTimeHz))); + time_id = std::clamp(time_id, 0, tokenizer.time_end - tokenizer.time_start - 1); + event.timestamp_tokens.clear(); + event.timestamp_tokens.push_back(static_cast(tokenizer.time_start + time_id)); + event.timestamp = static_cast(time_id) / static_cast(kTimeHz); +} + +std::vector build_overlap_prefix_events( + const std::vector & stitched_events, + const SheetSage2Tokenizer & tokenizer, + const SheetSage2Window & window, + int64_t & base_subbeat) { + constexpr double eps = 1.0e-4; + std::vector source_events; + for (const auto & event : stitched_events) { + const double t = event_time_value(event); + if (window.start - eps <= t && t < window.prefix_end - eps) { + source_events.push_back(event); + } + } + std::sort(source_events.begin(), source_events.end(), [](const auto & a, const auto & b) { + if (a.global_subbeat != b.global_subbeat) { + return a.global_subbeat < b.global_subbeat; + } + return a.time < b.time; + }); + auto first_beat = std::find_if(source_events.begin(), source_events.end(), [](const auto & event) { + return event.timestamp.has_value() || !event.rhythm_tokens.empty(); + }); + if (first_beat == source_events.end()) { + base_subbeat = 0; + return {}; + } + source_events.erase(source_events.begin(), first_beat); + base_subbeat = source_events.front().global_subbeat; + + std::optional> context_structure; + std::optional> context_key; + std::optional> context_chord; + std::optional> context_meter; + for (const auto & event : stitched_events) { + if (event.time > source_events.front().time + eps) { + continue; + } + if (!event.structure_tokens.empty()) { + context_structure = event.structure_tokens; + } + if (!event.key_tokens.empty()) { + context_key = event.key_tokens; + } + if (!event.chord_tokens.empty()) { + context_chord = event.chord_tokens; + } + for (const auto token : event.rhythm_tokens) { + if (token_type(tokenizer, token) == SheetSage2TokenType::Meter) { + context_meter = std::vector{token}; + break; + } + } + } + + std::vector prefix_events; + prefix_events.reserve(source_events.size()); + for (auto event : source_events) { + event.subbeat = std::max(0, event.global_subbeat - base_subbeat); + if (!event.timestamp_tokens.empty()) { + set_event_local_timestamp(event, tokenizer, static_cast(event.time) - window.start); + } + prefix_events.push_back(std::move(event)); + } + if (!prefix_events.empty()) { + auto & first = prefix_events.front(); + if (first.structure_tokens.empty() && context_structure.has_value()) { + first.structure_tokens = *context_structure; + } + if (first.key_tokens.empty() && context_key.has_value()) { + first.key_tokens = *context_key; + } + if (first.chord_tokens.empty() && context_chord.has_value()) { + first.chord_tokens = *context_chord; + } + const bool has_meter = std::any_of(first.rhythm_tokens.begin(), first.rhythm_tokens.end(), [&](int32_t token) { + return token_type(tokenizer, token) == SheetSage2TokenType::Meter; + }); + const bool has_eighth = std::any_of(first.rhythm_tokens.begin(), first.rhythm_tokens.end(), [&](int32_t token) { + return token_type(tokenizer, token) == SheetSage2TokenType::Eighth; + }); + if (!has_meter && has_eighth && context_meter.has_value()) { + first.rhythm_tokens.insert(first.rhythm_tokens.begin(), context_meter->begin(), context_meter->end()); + } + } + return prefix_events; +} + +std::vector stitched_window_events( + const std::vector & decoded_events, + const std::function & time_lookup, + const SheetSage2Window & window, + double song_duration, + int64_t global_subbeat_base) { + constexpr double eps = 1.0e-4; + std::vector accepted; + for (auto event : decoded_events) { + const double local_time = time_lookup(event.subbeat); + const double abs_time = window.start + local_time; + if (abs_time < window.accept_start - eps || + abs_time >= window.accept_end - eps || + abs_time >= song_duration - eps) { + continue; + } + event.time = static_cast(std::clamp(abs_time, 0.0, song_duration)); + event.window_index = static_cast(window.index); + event.window_start = static_cast(window.start); + event.source_subbeat = event.subbeat; + event.global_subbeat = global_subbeat_base + event.subbeat; + if (event.timestamp.has_value()) { + event.timestamp = event.time; + } + event.note_end_times.clear(); + event.note_end_times.reserve(event.notes.size()); + for (const auto & note : event.notes) { + const double local_end = time_lookup(event.subbeat + note.duration_steps); + const double end_time = std::min(song_duration, std::max(abs_time + 0.04, window.start + local_end)); + event.note_end_times.push_back(static_cast(end_time)); + } + accepted.push_back(std::move(event)); + } + return accepted; +} + +std::vector prepare_normalized_mel( + const std::vector & waveform, + const SheetSage2Assets & assets, + size_t threads, + int64_t & mel_frames_out) { + if (static_cast(waveform.size()) <= assets.config.n_fft) { + throw std::runtime_error("SheetSage2 window audio is too short for the audio frontend"); + } + const int64_t target_samples = static_cast(waveform.size()); + const engine::audio::STFTConfig stft_config{ + assets.config.n_fft, + assets.config.hop_length, + assets.config.win_length, + true, + engine::audio::STFTPadMode::Reflect, + engine::audio::STFTFamily::Default, + }; + auto magnitude = engine::audio::STFT().compute_magnitude( + waveform, + assets.stft_window, + 1, + target_samples, + stft_config, + threads); + const int64_t freq_bins = assets.config.n_fft / 2 + 1; + const int64_t stft_frames = magnitude.shape.at(2); + const int64_t mel_frames = stft_frames - 1; + auto mel = engine::audio::MelFilterbank().compute_custom_sparse_from_magnitude( + magnitude.values, + 1, + freq_bins, + stft_frames, + mel_frames, + assets.mel_filterbank); + if (mel.shape.size() != 3 || mel.shape[1] != assets.config.mel_bins || mel.shape[2] != mel_frames) { + throw std::runtime_error("SheetSage2 mel frontend produced invalid shape"); + } + std::vector normalized(static_cast(mel_frames * assets.config.mel_bins), 0.0F); + for (int64_t t = 0; t < mel_frames; ++t) { + for (int64_t m = 0; m < assets.config.mel_bins; ++m) { + const float power = std::max(mel.values[static_cast(m * mel.shape[2] + t)], 1.0e-10F); + const float db = 10.0F * std::log10(power); + const float stdv = std::max(assets.mel_std[static_cast(m)], 1.0e-5F); + normalized[static_cast(t * assets.config.mel_bins + m)] = + (db - assets.mel_mean[static_cast(m)]) / stdv; + } + } + mel_frames_out = mel_frames; + return normalized; +} + +std::vector generate_tokens( + SheetSage2DecoderRuntime & decoder, + const std::vector & memory, + int64_t memory_steps, + const SheetSage2Tokenizer & tokenizer, + int64_t max_sequence_length, + const std::vector * prefix_tokens, + std::optional stop_time_seconds) { + auto ids = prefix_tokens == nullptr ? prompt_prefix(tokenizer) : *prefix_tokens; + if (ids.empty() || ids.front() != tokenizer.sos) { + throw std::runtime_error("SheetSage2 generation prefix must begin with sos"); + } + if (ids.back() == tokenizer.eos) { + ids.pop_back(); + } + const auto out_it = std::find(ids.begin(), ids.end(), static_cast(tokenizer.out)); + if (out_it == ids.end()) { + throw std::runtime_error("SheetSage2 generation prefix is missing out token"); + } + SheetSage2GrammarState state; + for (auto it = out_it + 1; it != ids.end(); ++it) { + update_grammar(tokenizer, state, *it); + } + if (max_sequence_length <= static_cast(ids.size())) { + throw std::runtime_error("SheetSage2 max_tokens must exceed the generation prefix length"); + } + decoder.reset_cached_decode(memory, memory_steps, max_sequence_length); + std::vector logits = decoder.prefill_cached_decode(ids); + for (int64_t step = static_cast(ids.size()); step < max_sequence_length; ++step) { + const int32_t token = select_next_token(logits, tokenizer, state); + ids.push_back(token); + bool finished = update_grammar(tokenizer, state, token); + if (!finished && stop_time_seconds.has_value() && + token >= tokenizer.time_start && token < tokenizer.time_end && + static_cast(token - tokenizer.time_start) / static_cast(kTimeHz) >= *stop_time_seconds) { + ids.push_back(static_cast(tokenizer.eos)); + finished = true; + } + if (finished) { + break; + } + logits = decoder.decode_cached_step(token); + } + if (ids.back() != tokenizer.eos) { + ids.push_back(static_cast(tokenizer.eos)); + } + return ids; +} + +engine::audio::AudioTensor load_mel_filterbank(const assets::TensorSource & source, int64_t n_fft, int64_t mel_bins) { + const int64_t freq_bins = n_fft / 2 + 1; + const auto values = source.require_f32("feature_extractor.mel_scale.fb", {freq_bins, mel_bins}); + engine::audio::AudioTensor filterbank; + filterbank.shape = {mel_bins, freq_bins}; + filterbank.values.assign(static_cast(mel_bins * freq_bins), 0.0F); + for (int64_t f = 0; f < freq_bins; ++f) { + for (int64_t m = 0; m < mel_bins; ++m) { + filterbank.values[static_cast(m * freq_bins + f)] = + values[static_cast(f * mel_bins + m)]; + } + } + return filterbank; +} + +std::unique_ptr create_sheetsage2_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options, + std::shared_ptr assets, + std::shared_ptr contract) { + return std::make_unique( + task, + options, + std::move(assets), + std::move(contract)); +} + +} // namespace + +std::shared_ptr load_sheetsage2_assets(const std::filesystem::path & model_path) { + auto assets = std::make_shared(); + assets->resources = model_spec::load_resource_bundle_for_family(model_path, kFamily); + assets->config = parse_decoder_config(assets->resources); + assets->weights = assets->resources.open_tensor_source("weights"); + assets::require_tensor_shape(*assets->weights, "token_embedding.weight", {assets->config.vocab_size, assets->config.hidden_size}); + assets::require_tensor_shape( + *assets->weights, + "encoder_projection.weight", + {assets->config.hidden_size, assets->config.encoder_hidden_size}); + assets::require_tensor_shape(*assets->weights, "feature_extractor.mel_mean", {assets->config.mel_bins}); + assets::require_tensor_shape(*assets->weights, "feature_extractor.mel_std", {assets->config.mel_bins}); + assets::require_tensor_shape( + *assets->weights, + "feature_extractor.mel_scale.fb", + {assets->config.n_fft / 2 + 1, assets->config.mel_bins}); + assets::require_tensor_shape(*assets->weights, "subsampling_module.0.convnext_layers.0.depthwise_block.1.weight", {128, 1, 7}); + assets::require_tensor_shape( + *assets->weights, + "layers.0.attn.query_proj.weight", + {assets->config.encoder_hidden_size, assets->config.encoder_hidden_size}); + assets->mel_mean = assets->weights->require_f32("feature_extractor.mel_mean", {assets->config.mel_bins}); + assets->mel_std = assets->weights->require_f32("feature_extractor.mel_std", {assets->config.mel_bins}); + assets->stft_window = assets->weights->require_f32("feature_extractor.spectrogram.window", {assets->config.win_length}); + assets->mel_filterbank = engine::audio::MelFilterbank().prepare_sparse( + load_mel_filterbank(*assets->weights, assets->config.n_fft, assets->config.mel_bins)); + return assets; +} + +SheetSage2Session::SheetSage2Session( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract) + : RuntimeSessionBase(std::move(options)), + task_(std::move(task)), + assets_(require_assets(std::move(assets))), + contract_(require_contract(std::move(contract))), + encoder_( + assets_->weights, + execution_context(), + assets_->config, + decoder_options_from_session_options(RuntimeSessionBase::options())), + decoder_( + assets_->weights, + execution_context(), + assets_->config, + decoder_options_from_session_options(RuntimeSessionBase::options())) { + if (task_.task != runtime::VoiceTaskKind::Midi || task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("SheetSage2 supports only offline midi"); + } + runtime::validate_spec_backed_session_options(RuntimeSessionBase::options(), *contract_, kFamily, "SheetSage2"); +} + +SheetSage2Session::~SheetSage2Session() = default; + +std::string SheetSage2Session::family() const { + return kFamily; +} + +runtime::VoiceTaskKind SheetSage2Session::task_kind() const { + return task_.task; +} + +runtime::RunMode SheetSage2Session::run_mode() const { + return task_.mode; +} + +void SheetSage2Session::prepare(const runtime::SessionPreparationRequest & request) { + runtime::validate_spec_backed_request_options(request.options, *contract_, "SheetSage2"); + mark_prepared(); +} + +runtime::TaskResult SheetSage2Session::run(const runtime::TaskRequest & request) { + require_prepared("SheetSage2 run"); + runtime::validate_spec_backed_request_options(request.options, *contract_, "SheetSage2"); + if (!request.audio_input.has_value()) { + throw std::runtime_error("SheetSage2 run() requires audio_input"); + } + const auto total_start = Clock::now(); + const int64_t max_tokens = runtime::parse_positive_i64_option( + request.options, + {"max_tokens", "sheetsage2.max_tokens"}, + assets_->config.max_position_embeddings); + if (max_tokens > assets_->config.max_position_embeddings) { + throw std::runtime_error("SheetSage2 max_tokens exceeds the decoder context"); + } + const SheetSage2Tokenizer tokenizer(assets_->config.input_audio_length_samples / assets_->config.sampling_rate); + if (tokenizer.vocab_size != assets_->config.vocab_size) { + throw std::runtime_error("SheetSage2 tokenizer vocabulary does not match the loaded model"); + } + + const auto frontend_start = Clock::now(); + const auto waveform = audio_frontend_.prepare( + request.audio_input->samples, + request.audio_input->sample_rate, + request.audio_input->channels, + static_cast(assets_->config.sampling_rate), + std::max(1, RuntimeSessionBase::options().backend.threads)); + if (static_cast(waveform.size()) <= assets_->config.n_fft) { + throw std::runtime_error("SheetSage2 input audio is too short for the audio frontend"); + } + engine::debug::timing_log_scalar("sheetsage2.frontend_ms", engine::debug::elapsed_ms(frontend_start, Clock::now())); + + const auto windows = sliding_window_plan( + static_cast(waveform.size()), + assets_->config.sampling_rate, + assets_->config.input_audio_length_samples); + engine::debug::trace_log_scalar("sheetsage2.windows", static_cast(windows.size())); + engine::debug::trace_log_scalar( + "sheetsage2.duration_seconds", + static_cast(waveform.size()) / static_cast(assets_->config.sampling_rate)); + + double frontend_windows_ms = 0.0; + double encoder_ms = 0.0; + double generation_ms = 0.0; + std::vector stitched_events; + std::vector last_ids; + int64_t last_memory_steps = 0; + const double window_seconds = + static_cast(assets_->config.input_audio_length_samples) / static_cast(assets_->config.sampling_rate); + + for (const auto & window : windows) { + const auto window_frontend_start = Clock::now(); + const auto window_audio = slice_window_audio(waveform, window.start_sample, assets_->config.input_audio_length_samples); + int64_t mel_frames = 0; + const auto normalized_mel = prepare_normalized_mel( + window_audio, + *assets_, + static_cast(std::max(1, RuntimeSessionBase::options().backend.threads)), + mel_frames); + frontend_windows_ms += engine::debug::elapsed_ms(window_frontend_start, Clock::now()); + + const auto encoder_start = Clock::now(); + const auto mixed = encoder_.encode_mel(normalized_mel, mel_frames); + const int64_t memory_steps = static_cast(mixed.size()) / assets_->config.encoder_hidden_size; + encoder_ms += engine::debug::elapsed_ms(encoder_start, Clock::now()); + last_memory_steps = memory_steps; + + int64_t base_subbeat = 0; + std::vector prefix_tokens; + if (window.index > 0) { + const auto prefix_events = build_overlap_prefix_events(stitched_events, tokenizer, window, base_subbeat); + if (!prefix_events.empty()) { + prefix_tokens = encode_events(tokenizer, prefix_events, false); + } + } + + const auto generation_start = Clock::now(); + last_ids = generate_tokens( + decoder_, + mixed, + memory_steps, + tokenizer, + max_tokens, + prefix_tokens.empty() ? nullptr : &prefix_tokens, + window.generation_stop); + generation_ms += engine::debug::elapsed_ms(generation_start, Clock::now()); + + const auto decoded = decode_events(tokenizer, last_ids); + const auto time_lookup = make_event_time_lookup(decoded, window_seconds); + auto accepted = stitched_window_events( + decoded, + time_lookup, + window, + static_cast(waveform.size()) / static_cast(assets_->config.sampling_rate), + base_subbeat); + stitched_events.insert( + stitched_events.end(), + std::make_move_iterator(accepted.begin()), + std::make_move_iterator(accepted.end())); + engine::debug::trace_log_scalar("sheetsage2.window.index", window.index); + engine::debug::trace_log_scalar("sheetsage2.window.prefix_tokens", static_cast(prefix_tokens.size())); + engine::debug::trace_log_scalar("sheetsage2.window.tokens", static_cast(last_ids.size())); + } + encoder_.release_runtime_graphs(); + decoder_.release_runtime_graphs(); + + std::sort(stitched_events.begin(), stitched_events.end(), [](const auto & a, const auto & b) { + if (a.time != b.time) { + return a.time < b.time; + } + return a.global_subbeat < b.global_subbeat; + }); + engine::debug::timing_log_scalar("sheetsage2.frontend_windows_ms", frontend_windows_ms); + engine::debug::timing_log_scalar("sheetsage2.encoder.total_ms", encoder_ms); + engine::debug::timing_log_scalar("sheetsage2.decoder.generate_ms", generation_ms); + + const auto post_start = Clock::now(); + const auto abc = events_to_abc( + stitched_events, + static_cast(waveform.size()) / static_cast(assets_->config.sampling_rate)); + const auto json = events_json(stitched_events); + engine::debug::timing_log_scalar("sheetsage2.postprocess_ms", engine::debug::elapsed_ms(post_start, Clock::now())); + + runtime::TaskResult result; + result.text_output = runtime::Transcript{abc, "abc"}; + result.artifact_output = runtime::make_text_artifact( + runtime::ArtifactKind::Custom, + "score", + abc, + { + {"mime", "text/vnd.abc"}, + {"format", "abc"}, + {"extension", "abc"}, + {"tokens", std::to_string(last_ids.size())}, + {"memory_steps", std::to_string(last_memory_steps)}, + {"windows", std::to_string(windows.size())}, + }); + result.output_artifacts.push_back(runtime::make_text_artifact( + runtime::ArtifactKind::Custom, + "events", + json, + { + {"mime", "application/json"}, + {"format", "sheetsage2-events-json"}, + {"extension", "json"}, + })); + engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(total_start, Clock::now())); + return result; +} + +std::shared_ptr make_sheetsage2_loader() { + runtime::SpecBackedVoiceModelConfig config; + config.family = kFamily; + config.load_assets = load_sheetsage2_assets; + config.create_session = create_sheetsage2_session; + return runtime::make_spec_backed_voice_loader(std::move(config)); +} + +} // namespace engine::models::sheetsage diff --git a/src/models/yue2/ar_runtime.cpp b/src/models/yue2/ar_runtime.cpp new file mode 100644 index 00000000..6c006f47 --- /dev/null +++ b/src/models/yue2/ar_runtime.cpp @@ -0,0 +1,1043 @@ +#include "engine/models/yue2/ar_runtime.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/transformers/qwen_causal_decoder.h" +#include "engine/framework/modules/weight_binding.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::yue2 { +namespace { + +namespace binding = engine::modules::binding; +using Clock = std::chrono::steady_clock; + +constexpr int64_t kArDecodeChunkTokens = 5120; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct Yue2SamplerRange { + int64_t begin = 0; + int64_t end = 0; +}; + +struct Yue2SamplerScratch { + std::vector candidates; + std::vector kept; + std::vector weights; +}; + +engine::modules::QwenDecoderLayerWeights load_layer( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const Yue2ModelConfig & config, + assets::TensorStorageType storage_type, + int64_t layer) { + const std::string prefix = "model.layers." + std::to_string(layer); + engine::modules::QwenDecoderLayerWeights out; + out.input_norm = binding::norm_weight_from_source(store, source, prefix + ".input_layernorm", config.hidden_size); + out.self_attention.q_weight = store.load_tensor( + source, + prefix + ".self_attn.q_proj.weight", + storage_type, + {config.attention_heads * config.head_dim, config.hidden_size}); + out.self_attention.k_weight = store.load_tensor( + source, + prefix + ".self_attn.k_proj.weight", + storage_type, + {config.kv_heads * config.head_dim, config.hidden_size}); + out.self_attention.v_weight = store.load_tensor( + source, + prefix + ".self_attn.v_proj.weight", + storage_type, + {config.kv_heads * config.head_dim, config.hidden_size}); + out.self_attention.out_weight = store.load_tensor( + source, + prefix + ".self_attn.o_proj.weight", + storage_type, + {config.hidden_size, config.attention_heads * config.head_dim}); + out.q_norm = binding::norm_weight_from_source(store, source, prefix + ".self_attn.q_norm", config.head_dim); + out.k_norm = binding::norm_weight_from_source(store, source, prefix + ".self_attn.k_norm", config.head_dim); + out.post_norm = binding::norm_weight_from_source(store, source, prefix + ".post_attention_layernorm", config.hidden_size); + out.mlp.gate_proj = binding::linear_from_source( + store, + source, + prefix + ".mlp.gate_proj", + storage_type, + config.intermediate_size, + config.hidden_size, + false); + out.mlp.up_proj = binding::linear_from_source( + store, + source, + prefix + ".mlp.up_proj", + storage_type, + config.intermediate_size, + config.hidden_size, + false); + out.mlp.down_proj = binding::linear_from_source( + store, + source, + prefix + ".mlp.down_proj", + storage_type, + config.hidden_size, + config.intermediate_size, + false); + return out; +} + +engine::modules::QwenCausalDecodeRuntimeWeights load_prefix_weights( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const Yue2ModelConfig & config, + assets::TensorStorageType storage_type) { + engine::modules::QwenCausalDecodeRuntimeWeights weights; + weights.token_embedding = store.load_tensor( + source, + "model.embed_tokens.weight", + storage_type, + {config.vocab_size, config.hidden_size}); + weights.stack.layers.reserve(static_cast(config.layers)); + for (int64_t layer = 0; layer < config.layers; ++layer) { + weights.stack.layers.push_back(load_layer(store, source, config, storage_type, layer)); + } + return weights; +} + +void load_generation_weights( + engine::modules::QwenCausalDecodeRuntimeWeights & weights, + core::BackendWeightStore & store, + const assets::TensorSource & source, + const Yue2ModelConfig & config, + assets::TensorStorageType storage_type) { + weights.final_norm = binding::norm_weight_from_source(store, source, "model.norm", config.hidden_size); + weights.lm_head = binding::linear_from_source( + store, + source, + "lm_head", + storage_type, + config.vocab_size, + config.hidden_size, + false); +} + +engine::modules::QwenCausalDecodeRuntimeConfig make_runtime_config( + const Yue2ModelConfig & config, + core::BackendType backend_type, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + int64_t logits_size = 0) { + engine::modules::QwenCausalDecodeRuntimeConfig out; + out.trace_name = "yue2.ar"; + out.prefill_graph_arena_bytes = prefill_graph_arena_bytes; + out.decode_graph_arena_bytes = decode_graph_arena_bytes; + out.decoder.stack.hidden_size = config.hidden_size; + out.decoder.stack.layers = config.layers; + out.decoder.stack.num_attention_heads = config.attention_heads; + out.decoder.stack.num_key_value_heads = config.kv_heads; + out.decoder.stack.head_dim = config.head_dim; + out.decoder.stack.intermediate_size = config.intermediate_size; + out.decoder.stack.rms_norm_eps = config.rms_norm_eps; + out.decoder.stack.rope_theta = config.rope_theta; + out.decoder.stack.rope_type = GGML_ROPE_TYPE_NEOX; + out.decoder.stack.use_qk_norm = true; + out.decoder.stack.runtime.attention.prefill_mode = engine::modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.attention.static_mode = engine::modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.static_cache.update_mode = engine::modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.decoder.stack.runtime.static_cache.set_rows_mode = + engine::modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + if (backend_type == core::BackendType::Cuda || backend_type == core::BackendType::Hip || + backend_type == core::BackendType::Vulkan) { + out.decoder.static_cache_type = GGML_TYPE_F16; + } + out.decoder.logits_size = logits_size > 0 ? logits_size : config.vocab_size; + out.decoder.logits_mode = engine::modules::QwenCausalDecoderLogitsMode::LastStep; + out.readback_round_type = GGML_TYPE_BF16; + return out; +} + +runtime::TransformerBatchedKVState make_cfg_batched_state( + const runtime::TransformerKVState & positive, + const runtime::TransformerKVState & negative) { + if (positive.layers.size() != negative.layers.size()) { + throw std::runtime_error("Yue2 CFG batched state layer count mismatch"); + } + const int64_t positive_steps = positive.layers.empty() ? 0 : positive.layers.front().valid_steps; + const int64_t negative_steps = negative.layers.empty() ? 0 : negative.layers.front().valid_steps; + const int64_t max_steps = std::max(positive_steps, negative_steps); + if (positive_steps <= 0 || negative_steps <= 0 || max_steps <= 0) { + throw std::runtime_error("Yue2 CFG batched state requires non-empty prefix states"); + } + runtime::TransformerBatchedKVState out; + out.batch_size = 2; + out.current_end = std::max(positive.current_end, negative.current_end); + out.current_end_by_batch = {positive.current_end, negative.current_end}; + out.valid_steps_by_batch = {positive_steps, negative_steps}; + out.layers.resize(positive.layers.size()); + for (size_t layer = 0; layer < positive.layers.size(); ++layer) { + const auto & pos = positive.layers[layer]; + const auto & neg = negative.layers[layer]; + if (pos.valid_steps != positive_steps || neg.valid_steps != negative_steps || + pos.key.size() != pos.value.size() || neg.key.size() != neg.value.size()) { + throw std::runtime_error("Yue2 CFG batched state source shape mismatch"); + } + if (pos.key.size() % static_cast(positive_steps) != 0 || + neg.key.size() % static_cast(negative_steps) != 0) { + throw std::runtime_error("Yue2 CFG batched state source step shape mismatch"); + } + const size_t row_elems = pos.key.size() / static_cast(positive_steps); + if (neg.key.size() / static_cast(negative_steps) != row_elems) { + throw std::runtime_error("Yue2 CFG batched state row size mismatch"); + } + auto & dst = out.layers[layer]; + dst.valid_steps = max_steps; + dst.key.resize((static_cast(positive_steps) + static_cast(negative_steps)) * row_elems); + dst.value.assign(dst.key.size(), 0.0F); + std::copy(pos.key.begin(), pos.key.end(), dst.key.begin()); + std::copy(pos.value.begin(), pos.value.end(), dst.value.begin()); + std::copy(neg.key.begin(), neg.key.end(), dst.key.begin() + static_cast(pos.key.size())); + std::copy(neg.value.begin(), neg.value.end(), dst.value.begin() + static_cast(pos.value.size())); + } + return out; +} + +void apply_repetition_penalty( + std::vector & logits, + const std::vector & emitted, + const Yue2ArSamplingWindow & window) { + const float penalty = window.sampling.repetition_penalty; + if (penalty == 1.0F || emitted.empty()) { + return; + } + const int64_t begin = std::max(0, static_cast(emitted.size()) - window.sampling.penalty_window); + for (int64_t i = begin; i < static_cast(emitted.size()); ++i) { + const int32_t token = emitted[static_cast(i)]; + if (token < 0 || token >= static_cast(logits.size())) { + continue; + } + float & value = logits[static_cast(token)]; + value = value < 0.0F ? value * penalty : value / penalty; + } +} + +int32_t sample_from_allowed_ranges( + std::vector & logits, + const std::array & ranges, + const Yue2SamplingConfig & sampling, + std::mt19937 & rng, + Yue2SamplerScratch & scratch, + const char * context) { + if (!(sampling.temperature > 0.0F) || !std::isfinite(sampling.temperature)) { + throw std::runtime_error(std::string(context) + " temperature must be finite and positive"); + } + if (sampling.top_p < 0.0F || sampling.top_p > 1.0F || !std::isfinite(sampling.top_p)) { + throw std::runtime_error(std::string(context) + " top_p must be finite and in [0, 1]"); + } + scratch.kept.clear(); + scratch.candidates.clear(); + const bool bounded_top_k = sampling.top_k > 0; + if (bounded_top_k) { + const size_t keep = static_cast(std::max(sampling.top_k, 1)); + auto worse_than = [&](int32_t lhs, int32_t rhs) { + const float lhs_score = logits[static_cast(lhs)]; + const float rhs_score = logits[static_cast(rhs)]; + if (lhs_score == rhs_score) { + return lhs > rhs; + } + return lhs_score > rhs_score; + }; + for (const auto & range : ranges) { + const int64_t begin = std::max(0, range.begin); + const int64_t end = std::min(range.end, static_cast(logits.size())); + for (int64_t token = begin; token < end; ++token) { + const float score = logits[static_cast(token)]; + if (!std::isfinite(score)) { + continue; + } + const auto token_i32 = static_cast(token); + if (scratch.candidates.size() < keep) { + scratch.candidates.push_back(token_i32); + std::push_heap(scratch.candidates.begin(), scratch.candidates.end(), worse_than); + } else { + const int32_t worst = scratch.candidates.front(); + const float worst_score = logits[static_cast(worst)]; + if (score > worst_score || (score == worst_score && token_i32 < worst)) { + std::pop_heap(scratch.candidates.begin(), scratch.candidates.end(), worse_than); + scratch.candidates.back() = token_i32; + std::push_heap(scratch.candidates.begin(), scratch.candidates.end(), worse_than); + } + } + } + } + scratch.kept = scratch.candidates; + } else { + for (const auto & range : ranges) { + const int64_t begin = std::max(0, range.begin); + const int64_t end = std::min(range.end, static_cast(logits.size())); + for (int64_t token = begin; token < end; ++token) { + if (std::isfinite(logits[static_cast(token)])) { + scratch.kept.push_back(static_cast(token)); + } + } + } + } + if (scratch.kept.empty()) { + throw std::runtime_error(std::string(context) + " sampler has no finite logits"); + } + + auto score_at = [&](int32_t token) { + return logits[static_cast(token)] / sampling.temperature; + }; + if (sampling.top_p < 1.0F) { + std::sort(scratch.kept.begin(), scratch.kept.end(), [&](int32_t lhs, int32_t rhs) { + const float lhs_score = score_at(lhs); + const float rhs_score = score_at(rhs); + if (lhs_score == rhs_score) { + return lhs < rhs; + } + return lhs_score < rhs_score; + }); + } else if (bounded_top_k) { + std::sort(scratch.kept.begin(), scratch.kept.end()); + } + float max_score = -std::numeric_limits::infinity(); + for (const int32_t token : scratch.kept) { + max_score = std::max(max_score, score_at(token)); + } + if (!std::isfinite(max_score)) { + throw std::runtime_error(std::string(context) + " sampler max score is invalid"); + } + scratch.weights.resize(scratch.kept.size()); + double total = 0.0; + for (size_t index = 0; index < scratch.kept.size(); ++index) { + scratch.weights[index] = std::exp(static_cast(score_at(scratch.kept[index]) - max_score)); + total += scratch.weights[index]; + } + if (!(total > 0.0) || !std::isfinite(total)) { + throw std::runtime_error(std::string(context) + " probability mass is invalid"); + } + + if (sampling.top_p < 1.0F) { + double cumulative = 0.0; + size_t kept_count = 0; + const double remove_mass = 1.0 - static_cast(sampling.top_p); + const size_t protected_from = scratch.kept.size() - 1; + for (size_t index = 0; index < scratch.kept.size(); ++index) { + cumulative += scratch.weights[index] / total; + if (index < protected_from && cumulative <= remove_mass) { + continue; + } + scratch.kept[kept_count] = scratch.kept[index]; + scratch.weights[kept_count] = scratch.weights[index]; + ++kept_count; + } + scratch.kept.resize(kept_count); + scratch.weights.resize(kept_count); + } + + std::discrete_distribution distribution(scratch.weights.begin(), scratch.weights.end()); + return scratch.kept[distribution(rng)]; +} + +int32_t sample_token( + std::vector & logits, + const std::vector & emitted, + const Yue2ArSamplingWindow & window, + std::mt19937 & rng, + Yue2SamplerScratch & scratch) { + if (window.begin < 0 || window.end <= window.begin || + window.stop_token < 0 || static_cast(logits.size()) <= window.stop_token || + static_cast(logits.size()) < window.end) { + throw std::runtime_error("Yue2 AR sampling window is invalid"); + } + apply_repetition_penalty(logits, emitted, window); + return sample_from_allowed_ranges( + logits, + { + Yue2SamplerRange{window.begin, window.end}, + Yue2SamplerRange{ + static_cast(emitted.size()) >= window.min_tokens ? window.stop_token : 0, + static_cast(emitted.size()) >= window.min_tokens ? window.stop_token + 1 : 0, + }, + }, + window.sampling, + rng, + scratch, + "Yue2 AR"); +} + +int32_t compact_semantic_index(int32_t token) { + if (token == kMusicEndToken) { + return 0; + } + if (token >= kCodecOffset && token < kCodecOffset + kCodecSize) { + return token - kCodecOffset + 1; + } + return -1; +} + +int32_t sample_semantic_token( + std::vector & logits, + const std::vector & emitted, + const Yue2ArSamplingWindow & window, + std::mt19937 & rng, + Yue2SamplerScratch & scratch) { + if (static_cast(logits.size()) != kCodecSize + 1 || + window.begin != kCodecOffset || + window.end != kCodecOffset + kCodecSize || + window.stop_token != kMusicEndToken) { + throw std::runtime_error("Yue2 compact semantic logits shape mismatch"); + } + const float penalty = window.sampling.repetition_penalty; + if (penalty != 1.0F && !emitted.empty()) { + const int64_t begin = std::max(0, static_cast(emitted.size()) - window.sampling.penalty_window); + for (int64_t i = begin; i < static_cast(emitted.size()); ++i) { + const int32_t index = compact_semantic_index(emitted[static_cast(i)]); + if (index < 0 || index >= static_cast(logits.size())) { + continue; + } + float & value = logits[static_cast(index)]; + value = value < 0.0F ? value * penalty : value / penalty; + } + } + const int32_t index = sample_from_allowed_ranges( + logits, + { + Yue2SamplerRange{1, static_cast(logits.size())}, + Yue2SamplerRange{ + static_cast(emitted.size()) >= window.min_tokens ? 0 : 0, + static_cast(emitted.size()) >= window.min_tokens ? 1 : 0, + }, + }, + window.sampling, + rng, + scratch, + "Yue2 semantic AR"); + return index == 0 ? kMusicEndToken : kCodecOffset + index - 1; +} + +core::TensorValue view_linear_rows( + ggml_context * ctx, + const core::TensorValue & weight, + int64_t row_offset, + int64_t rows, + int64_t cols, + const char * label) { + if (weight.shape.rank != 2 || + weight.shape.dims[0] < row_offset + rows || + weight.shape.dims[1] != cols) { + throw std::runtime_error(std::string("Yue2 ") + label + " weight view is invalid"); + } + const size_t row_stride = weight.tensor->nb[1]; + const size_t byte_offset = static_cast(row_offset) * row_stride; + return core::wrap_tensor( + ggml_view_2d(ctx, weight.tensor, cols, rows, row_stride, byte_offset), + core::TensorShape::from_dims({rows, cols}), + weight.type); +} + +} // namespace + +struct Yue2ArRuntime::Impl { + Impl( + core::ExecutionContext & execution, + std::shared_ptr assets, + assets::TensorStorageType weight_type, + size_t weight_context_bytes, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes) + : execution(execution), + assets(std::move(assets)), + weight_type(weight_type) { + const auto total_start = Clock::now(); + if (!this->assets) { + throw std::runtime_error("Yue2 AR runtime requires assets"); + } + store = std::make_shared( + execution.backend(), + execution.backend_type(), + "yue2.ar.weights", + weight_context_bytes); + const auto & config = this->assets->config.model; + const auto & source = *this->assets->model_weights; + const auto load_start = Clock::now(); + runtime_weights = load_prefix_weights(*store, source, config, weight_type); + engine::debug::timing_log_scalar("yue2.ar.weights_load_ms", engine::debug::elapsed_ms(load_start)); + const auto upload_start = Clock::now(); + store->upload(); + engine::debug::timing_log_scalar("yue2.ar.prefix_weights_upload_ms", engine::debug::elapsed_ms(upload_start)); + engine::debug::timing_log_scalar("yue2.ar.weights_upload_ms", engine::debug::elapsed_ms(upload_start)); + runtime_config = make_runtime_config( + config, + execution.backend_type(), + prefill_graph_arena_bytes, + decode_graph_arena_bytes); + abc_runtime_config = make_runtime_config( + config, + execution.backend_type(), + prefill_graph_arena_bytes, + decode_graph_arena_bytes, + kAbcEndToken + 1); + semantic_runtime_config = make_runtime_config( + config, + execution.backend_type(), + prefill_graph_arena_bytes, + decode_graph_arena_bytes, + kCodecSize + 1); + engine::debug::timing_log_scalar("yue2.ar.init_total_ms", engine::debug::elapsed_ms(total_start)); + } + + struct PrefixStateGraph { + PrefixStateGraph(Impl & owner, int64_t steps) + : owner(&owner), + steps(steps) { + const auto total_start = Clock::now(); + const auto & config = owner.assets->config.model; + ggml_init_params params{owner.runtime_config.prefill_graph_arena_bytes, nullptr, true}; + ctx.reset(ggml_init(params)); + ggml_init_params state_params{ + ggml_tensor_overhead() * static_cast(config.layers * 2), + nullptr, + true}; + state_ctx.reset(ggml_init(state_params)); + if (ctx == nullptr || state_ctx == nullptr) { + throw std::runtime_error("failed to initialize Yue2 AR prefix-state graph context"); + } + core::ModuleBuildContext build{ctx.get(), "yue2.ar.prefix_state", owner.execution.backend_type()}; + core::ModuleBuildContext state_build{state_ctx.get(), "yue2.ar.prefix_state.cache", owner.execution.backend_type()}; + const auto build_start = Clock::now(); + input = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, steps); + positions = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, steps); + attention_mask = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, steps, steps, 1, 1); + auto ids = core::wrap_tensor(input, core::TensorShape::from_dims({steps}), GGML_TYPE_I32); + auto pos = core::wrap_tensor(positions, core::TensorShape::from_dims({steps}), GGML_TYPE_I32); + auto mask = core::wrap_tensor(attention_mask, core::TensorShape::from_dims({1, 1, steps, steps}), GGML_TYPE_F16); + auto x = engine::modules::EmbeddingModule({config.vocab_size, config.hidden_size}) + .build(build, ids, owner.runtime_weights.token_embedding); + x = core::reshape_tensor(build, x, core::TensorShape::from_dims({1, steps, config.hidden_size})); + auto stack = engine::modules::QwenDecoderStackModule(owner.runtime_config.decoder.stack) + .build(build, x, pos, owner.runtime_weights.stack, std::nullopt, mask); + keys.reserve(stack.state.layers.size()); + values.reserve(stack.state.layers.size()); + key_values.reserve(stack.state.layers.size()); + value_values.reserve(stack.state.layers.size()); + for (const auto & layer : stack.state.layers) { + if (!layer.key.has_value() || !layer.value.has_value()) { + throw std::runtime_error("Yue2 AR prefix-state graph did not return K/V state"); + } + const bool skip_bf16_round = owner.execution.backend_type() == core::BackendType::Metal; + auto key_value = core::wrap_tensor( + skip_bf16_round ? layer.key->tensor : ggml_round_bf16(ctx.get(), layer.key->tensor), + layer.key->shape, + GGML_TYPE_F32); + auto value_value = core::wrap_tensor( + skip_bf16_round ? layer.value->tensor : ggml_round_bf16(ctx.get(), layer.value->tensor), + layer.value->shape, + GGML_TYPE_F32); + key_value = core::wrap_tensor( + ggml_cast(ctx.get(), key_value.tensor, GGML_TYPE_F16), + key_value.shape, + GGML_TYPE_F16); + value_value = core::wrap_tensor( + ggml_cast(ctx.get(), value_value.tensor, GGML_TYPE_F16), + value_value.shape, + GGML_TYPE_F16); + auto * key = ggml_cpy(ctx.get(), key_value.tensor, ggml_dup_tensor(ctx.get(), key_value.tensor)); + auto * value = ggml_cpy(ctx.get(), value_value.tensor, ggml_dup_tensor(ctx.get(), value_value.tensor)); + ggml_set_output(key); + ggml_set_output(value); + keys.push_back(key); + values.push_back(value); + key_values.push_back(core::make_tensor(state_build, GGML_TYPE_F16, layer.key->shape)); + value_values.push_back(core::make_tensor(state_build, GGML_TYPE_F16, layer.value->shape)); + } + graph = ggml_new_graph_custom(ctx.get(), 65536, false); + for (auto * key : keys) { + ggml_build_forward_expand(graph, key); + } + for (auto * value : values) { + ggml_build_forward_expand(graph, value); + } + gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(owner.execution.backend())); + if (gallocr == nullptr || + !ggml_gallocr_reserve(gallocr, graph) || + !ggml_gallocr_alloc_graph(gallocr, graph)) { + throw std::runtime_error("failed to allocate Yue2 AR prefix-state graph"); + } + state_buffer = ggml_backend_alloc_ctx_tensors(state_ctx.get(), owner.execution.backend()); + if (state_buffer == nullptr) { + throw std::runtime_error("failed to allocate Yue2 AR prefix-state cache"); + } + position_values = engine::modules::qwen_position_ids(steps); + mask_values = engine::modules::qwen_causal_prefill_mask_values(1, steps); + ggml_backend_tensor_set(positions, position_values.data(), 0, position_values.size() * sizeof(int32_t)); + ggml_backend_tensor_set(attention_mask, mask_values.data(), 0, mask_values.size() * sizeof(ggml_fp16_t)); + engine::debug::timing_log_scalar("yue2.ar.prefix_state.graph.build_ms", engine::debug::elapsed_ms(build_start)); + engine::debug::timing_log_scalar("yue2.ar.prefix_state.graph.total_ms", engine::debug::elapsed_ms(total_start)); + } + + ~PrefixStateGraph() { + core::release_backend_graph_resources(owner->execution.backend(), graph); + if (gallocr != nullptr) { + ggml_gallocr_free(gallocr); + } + if (state_buffer != nullptr) { + ggml_backend_buffer_free(state_buffer); + } + } + + bool matches(int64_t s) const noexcept { + return steps == s; + } + + runtime::TransformerKVState run(const std::vector & tokens) { + compute(tokens); + const auto read_start = Clock::now(); + runtime::TransformerKVState out; + out.current_end = steps; + out.layers.resize(keys.size()); + for (size_t layer = 0; layer < keys.size(); ++layer) { + auto & state = out.layers[layer]; + state.valid_steps = steps; + core::read_tensor_f32_into(keys[layer], state.key); + core::read_tensor_f32_into(values[layer], state.value); + core::round_f32_to_bf16_in_place(state.key); + core::round_f32_to_bf16_in_place(state.value); + } + engine::debug::timing_log_scalar("yue2.ar.prefix_state.output_read_ms", engine::debug::elapsed_ms(read_start)); + return out; + } + + Yue2ArDevicePrefixState run_device(const std::vector & tokens) { + compute(tokens); + const auto copy_start = Clock::now(); + for (size_t layer = 0; layer < keys.size(); ++layer) { + ggml_backend_tensor_copy(keys[layer], key_values[layer].tensor); + ggml_backend_tensor_copy(values[layer], value_values[layer].tensor); + } + engine::debug::timing_log_scalar("yue2.ar.prefix_state.device_copy_ms", engine::debug::elapsed_ms(copy_start)); + Yue2ArDevicePrefixState out; + out.current_end = steps; + out.keys = key_values; + out.values = value_values; + return out; + } + + void compute(const std::vector & tokens) { + if (static_cast(tokens.size()) != steps) { + throw std::runtime_error("Yue2 AR prefix-state token size mismatch"); + } + const auto upload_start = Clock::now(); + ggml_backend_tensor_set(input, tokens.data(), 0, tokens.size() * sizeof(int32_t)); + ggml_backend_tensor_set(positions, position_values.data(), 0, position_values.size() * sizeof(int32_t)); + ggml_backend_tensor_set(attention_mask, mask_values.data(), 0, mask_values.size() * sizeof(ggml_fp16_t)); + engine::debug::timing_log_scalar("yue2.ar.prefix_state.input_upload_ms", engine::debug::elapsed_ms(upload_start)); + const auto compute_start = Clock::now(); + const auto status = core::compute_backend_graph(owner->execution.backend(), graph, nullptr, "yue2.ar.prefix_state"); + ggml_backend_synchronize(owner->execution.backend()); + engine::debug::timing_log_scalar("yue2.ar.prefix_state.graph_compute_ms", engine::debug::elapsed_ms(compute_start)); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Yue2 AR prefix-state graph compute failed"); + } + } + + Impl * owner = nullptr; + int64_t steps = 0; + std::unique_ptr ctx; + std::unique_ptr state_ctx; + ggml_cgraph * graph = nullptr; + ggml_gallocr_t gallocr = nullptr; + ggml_backend_buffer_t state_buffer = nullptr; + ggml_tensor * input = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * attention_mask = nullptr; + std::vector keys; + std::vector values; + std::vector key_values; + std::vector value_values; + std::vector position_values; + std::vector mask_values; + }; + + std::vector generate( + const std::vector & prefix, + const Yue2ArSamplingWindow & window, + uint64_t seed) { + if (prefix.empty()) { + throw std::runtime_error("Yue2 AR prefix must not be empty"); + } + const bool compact_semantic = is_semantic_window(window); + const bool compact_abc = is_abc_window(window); + ensure_generation_runtime(false, compact_semantic, compact_abc); + auto & active_runtime = compact_semantic ? semantic_runtime : (compact_abc ? abc_runtime : runtime); + const auto total_start = Clock::now(); + engine::debug::timing_log_scalar("yue2.ar.generate.prefix_tokens", prefix.size()); + engine::debug::timing_log_scalar("yue2.ar.generate.max_tokens", window.max_tokens); + auto cache_steps_for = [](int64_t prefix_tokens, int64_t remaining_tokens) { + return prefix_tokens + std::min(remaining_tokens, kArDecodeChunkTokens); + }; + const auto prefill_start = Clock::now(); + auto prefill = active_runtime->prefill_tokens_into_decode_cache( + prefix, + cache_steps_for(static_cast(prefix.size()), window.max_tokens)); + engine::debug::timing_log_scalar("yue2.ar.generate.prefill_ms", engine::debug::elapsed_ms(prefill_start)); + double start_decode_ms = 0.0; + std::vector emitted; + emitted.reserve(static_cast(window.max_tokens)); + std::mt19937 rng(static_cast(seed)); + Yue2SamplerScratch scratch; + engine::modules::QwenCausalDecodeStepResult decode_result; + decode_result.logits = std::move(prefill.logits); + decode_result.hidden = std::move(prefill.hidden); + double sample_ms = 0.0; + double decode_ms = 0.0; + double refill_prefill_ms = 0.0; + int64_t refill_count = 0; + for (int64_t step = 0; step < window.max_tokens; ++step) { + const auto sample_start = Clock::now(); + const int32_t token = compact_semantic ? + sample_semantic_token(decode_result.logits, emitted, window, rng, scratch) : + sample_token(decode_result.logits, emitted, window, rng, scratch); + sample_ms += engine::debug::elapsed_ms(sample_start); + if (token == window.stop_token) { + engine::debug::timing_log_scalar("yue2.ar.generate.start_decode_ms", start_decode_ms); + engine::debug::timing_log_scalar("yue2.ar.generate.sample_ms", sample_ms); + engine::debug::timing_log_scalar("yue2.ar.generate.decode_ms", decode_ms); + engine::debug::timing_log_scalar("yue2.ar.generate.refill_prefill_ms", refill_prefill_ms); + engine::debug::timing_log_scalar("yue2.ar.generate.refill_count", refill_count); + engine::debug::timing_log_scalar("yue2.ar.generate.emitted_tokens", emitted.size()); + engine::debug::timing_log_scalar("yue2.ar.generate.total_ms", engine::debug::elapsed_ms(total_start)); + return emitted; + } + emitted.push_back(token); + if (static_cast(emitted.size()) >= window.max_tokens) { + break; + } + if (active_runtime->decode_valid_steps() >= active_runtime->decode_cache_steps()) { + std::vector refill_prefix; + refill_prefix.reserve(prefix.size() + emitted.size()); + refill_prefix.insert(refill_prefix.end(), prefix.begin(), prefix.end()); + refill_prefix.insert(refill_prefix.end(), emitted.begin(), emitted.end()); + const auto refill_prefill_start = Clock::now(); + const int64_t remaining = window.max_tokens - static_cast(emitted.size()); + auto refill = active_runtime->prefill_tokens_into_decode_cache( + refill_prefix, + cache_steps_for(static_cast(refill_prefix.size()), remaining)); + refill_prefill_ms += engine::debug::elapsed_ms(refill_prefill_start); + decode_result.logits = std::move(refill.logits); + decode_result.hidden = std::move(refill.hidden); + ++refill_count; + continue; + } + const auto decode_start = Clock::now(); + active_runtime->decode_token_into(token, decode_result); + decode_ms += engine::debug::elapsed_ms(decode_start); + } + engine::debug::timing_log_scalar("yue2.ar.generate.start_decode_ms", start_decode_ms); + engine::debug::timing_log_scalar("yue2.ar.generate.sample_ms", sample_ms); + engine::debug::timing_log_scalar("yue2.ar.generate.decode_ms", decode_ms); + engine::debug::timing_log_scalar("yue2.ar.generate.refill_prefill_ms", refill_prefill_ms); + engine::debug::timing_log_scalar("yue2.ar.generate.refill_count", refill_count); + engine::debug::timing_log_scalar("yue2.ar.generate.emitted_tokens", emitted.size()); + engine::debug::timing_log_scalar("yue2.ar.generate.total_ms", engine::debug::elapsed_ms(total_start)); + return emitted; + } + + std::vector generate_cfg( + const std::vector & positive_prefix, + const std::vector & negative_prefix, + const Yue2ArSamplingWindow & window, + float guidance_scale, + uint64_t seed) { + if (guidance_scale == 1.0F) { + return generate(positive_prefix, window, seed); + } + const bool compact_semantic = is_semantic_window(window); + const bool compact_abc = is_abc_window(window); + ensure_generation_runtime(false, compact_semantic, compact_abc); + auto & positive_runtime = compact_semantic ? semantic_runtime : (compact_abc ? abc_runtime : runtime); + const auto total_start = Clock::now(); + engine::debug::timing_log_scalar("yue2.ar.cfg.positive_prefix_tokens", positive_prefix.size()); + engine::debug::timing_log_scalar("yue2.ar.cfg.negative_prefix_tokens", negative_prefix.size()); + engine::debug::timing_log_scalar("yue2.ar.cfg.max_tokens", window.max_tokens); + engine::debug::timing_log_scalar("yue2.ar.cfg.guidance_scale", static_cast(guidance_scale)); + const auto positive_prefill_start = Clock::now(); + auto positive = positive_runtime->prefill_tokens(positive_prefix); + engine::debug::timing_log_scalar("yue2.ar.cfg.prefill_positive_ms", engine::debug::elapsed_ms(positive_prefill_start)); + const auto negative_prefill_start = Clock::now(); + auto negative = positive_runtime->prefill_tokens(negative_prefix); + engine::debug::timing_log_scalar("yue2.ar.cfg.prefill_negative_ms", engine::debug::elapsed_ms(negative_prefill_start)); + const auto start_decode_start = Clock::now(); + const int64_t cache_steps = + std::max( + static_cast(positive_prefix.size()), + static_cast(negative_prefix.size())) + + window.max_tokens; + positive_runtime->start_decode_tokens_batched( + make_cfg_batched_state(positive.state, negative.state), + cache_steps); + engine::debug::timing_log_scalar("yue2.ar.cfg.start_decode_ms", engine::debug::elapsed_ms(start_decode_start)); + std::vector emitted; + emitted.reserve(static_cast(window.max_tokens)); + std::mt19937 rng(static_cast(seed)); + Yue2SamplerScratch scratch; + std::vector logits(positive.logits.size(), 0.0F); + double guidance_ms = 0.0; + double sample_ms = 0.0; + double decode_batched_ms = 0.0; + for (int64_t step = 0; step < window.max_tokens; ++step) { + if (positive.logits.size() != negative.logits.size()) { + throw std::runtime_error("Yue2 CFG logits size mismatch"); + } + const auto guidance_start = Clock::now(); + for (size_t i = 0; i < logits.size(); ++i) { + logits[i] = negative.logits[i] + (positive.logits[i] - negative.logits[i]) * guidance_scale; + } + guidance_ms += engine::debug::elapsed_ms(guidance_start); + const auto sample_start = Clock::now(); + const int32_t token = compact_semantic ? + sample_semantic_token(logits, emitted, window, rng, scratch) : + sample_token(logits, emitted, window, rng, scratch); + sample_ms += engine::debug::elapsed_ms(sample_start); + if (token == window.stop_token) { + engine::debug::timing_log_scalar("yue2.ar.cfg.guidance_ms", guidance_ms); + engine::debug::timing_log_scalar("yue2.ar.cfg.sample_ms", sample_ms); + engine::debug::timing_log_scalar("yue2.ar.cfg.decode_batched_ms", decode_batched_ms); + engine::debug::timing_log_scalar("yue2.ar.cfg.emitted_tokens", emitted.size()); + engine::debug::timing_log_scalar("yue2.ar.cfg.total_ms", engine::debug::elapsed_ms(total_start)); + return emitted; + } + emitted.push_back(token); + if (static_cast(emitted.size()) >= window.max_tokens) { + break; + } + const auto decode_start = Clock::now(); + auto batched = positive_runtime->decode_tokens_batched({token, token}); + decode_batched_ms += engine::debug::elapsed_ms(decode_start); + if (batched.logits.size() % 2 != 0) { + throw std::runtime_error("Yue2 CFG batched logits shape mismatch"); + } + const size_t row = batched.logits.size() / 2; + positive.logits.assign(batched.logits.begin(), batched.logits.begin() + static_cast(row)); + negative.logits.assign(batched.logits.begin() + static_cast(row), batched.logits.end()); + } + engine::debug::timing_log_scalar("yue2.ar.cfg.guidance_ms", guidance_ms); + engine::debug::timing_log_scalar("yue2.ar.cfg.sample_ms", sample_ms); + engine::debug::timing_log_scalar("yue2.ar.cfg.decode_batched_ms", decode_batched_ms); + engine::debug::timing_log_scalar("yue2.ar.cfg.emitted_tokens", emitted.size()); + engine::debug::timing_log_scalar("yue2.ar.cfg.total_ms", engine::debug::elapsed_ms(total_start)); + return emitted; + } + + runtime::TransformerKVState prefill_state(const std::vector & tokens) { + const auto start = Clock::now(); + engine::debug::timing_log_scalar("yue2.ar.prefill_state.tokens", tokens.size()); + const int64_t steps = static_cast(tokens.size()); + if (!prefix_state_graph || !prefix_state_graph->matches(steps)) { + prefix_state_graph = std::make_unique(*this, steps); + } + auto state = prefix_state_graph->run(tokens); + engine::debug::timing_log_scalar("yue2.ar.prefill_state_ms", engine::debug::elapsed_ms(start)); + return state; + } + + Yue2ArDevicePrefixState prefill_device_state(const std::vector & tokens) { + const auto start = Clock::now(); + engine::debug::timing_log_scalar("yue2.ar.prefill_state.tokens", tokens.size()); + const int64_t steps = static_cast(tokens.size()); + if (!prefix_state_graph || !prefix_state_graph->matches(steps)) { + prefix_state_graph = std::make_unique(*this, steps); + } + auto state = prefix_state_graph->run_device(tokens); + engine::debug::timing_log_scalar("yue2.ar.prefill_state_ms", engine::debug::elapsed_ms(start)); + return state; + } + + static bool is_semantic_window(const Yue2ArSamplingWindow & window) noexcept { + return window.begin == kCodecOffset && + window.end == kCodecOffset + kCodecSize && + window.stop_token == kMusicEndToken; + } + + static bool is_abc_window(const Yue2ArSamplingWindow & window) noexcept { + return window.begin == 0 && + window.end == kEodToken && + window.stop_token == kAbcEndToken; + } + + void ensure_generation_runtime(bool require_negative, bool compact_semantic, bool compact_abc) { + auto & active_runtime = compact_semantic ? semantic_runtime : (compact_abc ? abc_runtime : runtime); + auto & active_negative_runtime = + compact_semantic ? semantic_negative_runtime : (compact_abc ? abc_negative_runtime : negative_runtime); + const auto & active_config = + compact_semantic ? semantic_runtime_config : (compact_abc ? abc_runtime_config : runtime_config); + if (active_runtime && (!require_negative || active_negative_runtime)) { + return; + } + const auto total_start = Clock::now(); + if (!generation_store) { + const auto & config = assets->config.model; + const auto & source = *assets->model_weights; + generation_store = std::make_shared( + execution.backend(), + execution.backend_type(), + "yue2.ar.generation.weights", + 64ull * 1024ull * 1024ull); + const auto bind_start = Clock::now(); + load_generation_weights(runtime_weights, *generation_store, source, config, weight_type); + engine::debug::timing_log_scalar( + "yue2.ar.generation_weights_bind_ms", + engine::debug::elapsed_ms(bind_start)); + const auto upload_start = Clock::now(); + generation_store->upload(); + engine::debug::timing_log_scalar( + "yue2.ar.generation_weights_upload_ms", + engine::debug::elapsed_ms(upload_start)); + ggml_init_params view_params{ggml_tensor_overhead() * 8, nullptr, true}; + generation_view_ctx.reset(ggml_init(view_params)); + if (generation_view_ctx == nullptr) { + throw std::runtime_error("failed to initialize Yue2 AR generation weight views"); + } + abc_runtime_weights = runtime_weights; + semantic_runtime_weights = runtime_weights; + if (!runtime_weights.lm_head.has_value()) { + throw std::runtime_error("Yue2 AR generation runtime requires lm_head"); + } + abc_runtime_weights.lm_head = engine::modules::LinearWeights{ + view_linear_rows( + generation_view_ctx.get(), + runtime_weights.lm_head->weight, + 0, + kAbcEndToken + 1, + config.hidden_size, + "ABC lm_head"), + std::nullopt, + }; + semantic_runtime_weights.lm_head = engine::modules::LinearWeights{ + view_linear_rows( + generation_view_ctx.get(), + runtime_weights.lm_head->weight, + kMusicEndToken, + kCodecSize + 1, + config.hidden_size, + "semantic lm_head"), + std::nullopt, + }; + } + const auto runtime_start = Clock::now(); + if (!active_runtime) { + active_runtime = std::make_unique( + execution, + active_config, + compact_semantic ? semantic_runtime_weights : (compact_abc ? abc_runtime_weights : runtime_weights)); + } + if (require_negative && !active_negative_runtime) { + active_negative_runtime = std::make_unique( + execution, + active_config, + compact_semantic ? semantic_runtime_weights : (compact_abc ? abc_runtime_weights : runtime_weights)); + } + engine::debug::timing_log_scalar("yue2.ar.runtime_build_ms", engine::debug::elapsed_ms(runtime_start)); + engine::debug::timing_log_scalar("yue2.ar.ensure_generation_ms", engine::debug::elapsed_ms(total_start)); + } + + core::ExecutionContext & execution; + std::shared_ptr assets; + std::shared_ptr store; + std::shared_ptr generation_store; + std::unique_ptr generation_view_ctx; + assets::TensorStorageType weight_type; + engine::modules::QwenCausalDecodeRuntimeConfig runtime_config; + engine::modules::QwenCausalDecodeRuntimeConfig abc_runtime_config; + engine::modules::QwenCausalDecodeRuntimeConfig semantic_runtime_config; + engine::modules::QwenCausalDecodeRuntimeWeights runtime_weights; + engine::modules::QwenCausalDecodeRuntimeWeights abc_runtime_weights; + engine::modules::QwenCausalDecodeRuntimeWeights semantic_runtime_weights; + std::unique_ptr runtime; + std::unique_ptr negative_runtime; + std::unique_ptr abc_runtime; + std::unique_ptr abc_negative_runtime; + std::unique_ptr semantic_runtime; + std::unique_ptr semantic_negative_runtime; + std::unique_ptr prefix_state_graph; +}; + +Yue2ArRuntime::Yue2ArRuntime( + core::ExecutionContext & execution, + std::shared_ptr assets, + assets::TensorStorageType weight_type, + size_t weight_context_bytes, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes) + : impl_(std::make_unique( + execution, + std::move(assets), + weight_type, + weight_context_bytes, + prefill_graph_arena_bytes, + decode_graph_arena_bytes)) {} + +Yue2ArRuntime::~Yue2ArRuntime() = default; + +std::vector Yue2ArRuntime::generate( + const std::vector & prefix, + const Yue2ArSamplingWindow & window, + uint64_t seed) { + return impl_->generate(prefix, window, seed); +} + +std::vector Yue2ArRuntime::generate_cfg( + const std::vector & positive_prefix, + const std::vector & negative_prefix, + const Yue2ArSamplingWindow & window, + float guidance_scale, + uint64_t seed) { + return impl_->generate_cfg(positive_prefix, negative_prefix, window, guidance_scale, seed); +} + +runtime::TransformerKVState Yue2ArRuntime::prefill_state(const std::vector & tokens) { + return impl_->prefill_state(tokens); +} + +Yue2ArDevicePrefixState Yue2ArRuntime::prefill_device_state(const std::vector & tokens) { + return impl_->prefill_device_state(tokens); +} + +void Yue2ArRuntime::release_runtime_graphs() { + impl_->prefix_state_graph.reset(); + if (impl_->runtime) { + impl_->runtime->release_runtime_graphs(); + } + if (impl_->negative_runtime) { + impl_->negative_runtime->release_runtime_graphs(); + } + if (impl_->abc_runtime) { + impl_->abc_runtime->release_runtime_graphs(); + } + if (impl_->abc_negative_runtime) { + impl_->abc_negative_runtime->release_runtime_graphs(); + } + if (impl_->semantic_runtime) { + impl_->semantic_runtime->release_runtime_graphs(); + } + if (impl_->semantic_negative_runtime) { + impl_->semantic_negative_runtime->release_runtime_graphs(); + } +} + +} // namespace engine::models::yue2 diff --git a/src/models/yue2/assets.cpp b/src/models/yue2/assets.cpp new file mode 100644 index 00000000..b4713696 --- /dev/null +++ b/src/models/yue2/assets.cpp @@ -0,0 +1,125 @@ +#include "engine/models/yue2/assets.h" + +#include "engine/framework/io/config.h" +#include "engine/framework/io/filesystem.h" +#include "engine/framework/io/json.h" + +#include +#include + +namespace engine::models::yue2 { +namespace json = engine::io::json; +namespace { + +Yue2SamplingConfig parse_sampling(const json::Value & value, Yue2SamplingConfig fallback) { + fallback.temperature = json::optional_f32(value, "temperature", fallback.temperature); + fallback.top_p = json::optional_f32(value, "top_p", fallback.top_p); + fallback.top_k = json::optional_i64(value, "top_k", fallback.top_k); + fallback.repetition_penalty = json::optional_f32(value, "repetition_penalty", fallback.repetition_penalty); + fallback.penalty_window = json::optional_i64(value, "penalty_window", fallback.penalty_window); + fallback.min_tokens = json::optional_i64(value, "min_tokens", fallback.min_tokens); + fallback.max_tokens = json::optional_i64(value, "max_tokens", fallback.max_tokens); + if (fallback.temperature < 0.0F || fallback.temperature > 5.0F || + fallback.top_p <= 0.0F || fallback.top_p > 1.0F || + fallback.top_k < 1 || + fallback.repetition_penalty <= 0.0F || + fallback.penalty_window < 1 || + fallback.min_tokens < 0 || + fallback.max_tokens < fallback.min_tokens) { + throw std::runtime_error("Yue2 sampling config is invalid"); + } + return fallback; +} + +Yue2ModelConfig parse_model_config(const std::filesystem::path & path) { + const auto root = json::parse_file(path); + Yue2ModelConfig out; + out.hidden_size = json::require_i64(root, "hidden_size"); + out.layers = json::require_i64(root, "num_hidden_layers"); + out.attention_heads = json::require_i64(root, "num_attention_heads"); + out.kv_heads = json::require_i64(root, "num_key_value_heads"); + out.head_dim = json::require_i64(root, "head_dim"); + out.intermediate_size = json::require_i64(root, "intermediate_size"); + out.vocab_size = json::require_i64(root, "vocab_size"); + out.max_position_embeddings = json::require_i64(root, "max_position_embeddings"); + out.latent_dim = json::optional_i64(root, "latent_dim", out.latent_dim); + out.max_latent_frames = json::optional_i64(root, "max_latent_frames", out.max_latent_frames); + out.rms_norm_eps = json::optional_f32(root, "rms_norm_eps", out.rms_norm_eps); + out.rope_theta = json::optional_f32(root, "rope_theta", out.rope_theta); + out.timestep_shift = json::optional_f32(root, "timestep_shift", out.timestep_shift); + engine::io::require_positive(out.hidden_size, "Yue2 hidden_size"); + engine::io::require_positive(out.layers, "Yue2 layers"); + engine::io::require_positive(out.attention_heads, "Yue2 attention heads"); + engine::io::require_positive(out.kv_heads, "Yue2 kv heads"); + engine::io::require_positive(out.head_dim, "Yue2 head_dim"); + engine::io::require_positive(out.intermediate_size, "Yue2 intermediate_size"); + engine::io::require_positive(out.vocab_size, "Yue2 vocab_size"); + engine::io::require_positive(out.max_position_embeddings, "Yue2 max_position_embeddings"); + if (out.attention_heads % out.kv_heads != 0) { + throw std::runtime_error("Yue2 attention heads must be divisible by kv heads"); + } + return out; +} + +Yue2VaeConfig parse_vae_config(const std::filesystem::path & path) { + const auto root = json::parse_file(path); + Yue2VaeConfig out; + out.sample_rate = static_cast(json::optional_i64(root, "sample_rate", out.sample_rate)); + out.channels = json::optional_i64(root, "audio_channels", out.channels); + out.latent_dim = json::optional_i64(root, "latent_dim", out.latent_dim); + out.downsampling_ratio = json::optional_i64(root, "downsampling_ratio", out.downsampling_ratio); + out.decode_core_frames = json::optional_i64(root, "decode_core_frames", out.decode_core_frames); + out.decode_halo_frames = json::optional_i64(root, "decode_halo_frames", out.decode_halo_frames); + engine::io::require_positive(out.sample_rate, "Yue2 VAE sample_rate"); + engine::io::require_positive(out.channels, "Yue2 VAE channels"); + engine::io::require_positive(out.latent_dim, "Yue2 VAE latent_dim"); + return out; +} + +Yue2GenerationConfig parse_generation_config(const std::filesystem::path & path) { + Yue2GenerationConfig out; + if (!engine::io::is_existing_file(path)) { + out.abc = Yue2SamplingConfig{0.7F, 0.9F, 30, 1.005F, 100, 32, 4096}; + return out; + } + const auto root = json::parse_file(path); + if (const auto * abc = root.find("abc")) { + out.abc = parse_sampling(*abc, Yue2SamplingConfig{0.7F, 0.9F, 30, 1.005F, 100, 32, 4096}); + } + if (const auto * semantic = root.find("semantic")) { + out.semantic = parse_sampling(*semantic, out.semantic); + } + out.ode_steps = json::optional_i64(root, "ode_steps", out.ode_steps); + out.context = json::optional_i64(root, "context", out.context); + if (out.ode_steps <= 0 || out.context != kContextTokens) { + throw std::runtime_error("Yue2 generation config requires positive midpoint steps and context=24576"); + } + return out; +} + +} // namespace + +std::shared_ptr load_yue2_assets(const std::filesystem::path & model_path) { + auto assets = std::make_shared(); + assets->model_root = assets::prepare_model_directory(model_path).model_root; + const auto sidecars = assets->model_root / "sidecars"; + assets->tiktoken_path = sidecars / "yue2-qwen.tiktoken"; + const auto model_config = sidecars / "yue2-model-config.json"; + const auto generation_config = sidecars / "yue2-generation-config.json"; + const auto vae_config = sidecars / "yue2-vae-config.json"; + if (!engine::io::is_existing_file(model_config)) { + throw std::runtime_error("missing Yue2 model config: " + model_config.string()); + } + if (!engine::io::is_existing_file(vae_config)) { + throw std::runtime_error("missing Yue2 VAE config: " + vae_config.string()); + } + if (!engine::io::is_existing_file(assets->tiktoken_path)) { + throw std::runtime_error("missing Yue2 tiktoken file: " + assets->tiktoken_path.string()); + } + assets->config.model = parse_model_config(model_config); + assets->config.vae = parse_vae_config(vae_config); + assets->config.generation = parse_generation_config(generation_config); + return assets; +} + +} // namespace engine::models::yue2 diff --git a/src/models/yue2/nar_runtime.cpp b/src/models/yue2/nar_runtime.cpp new file mode 100644 index 00000000..ed7f25e3 --- /dev/null +++ b/src/models/yue2/nar_runtime.cpp @@ -0,0 +1,754 @@ +#include "engine/models/yue2/nar_runtime.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/transformers/qwen_decoder.h" +#include "engine/framework/modules/weight_binding.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::yue2 { +namespace { + +namespace binding = engine::modules::binding; +using Clock = std::chrono::steady_clock; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct Yue2NarWeights { + std::shared_ptr store; + engine::modules::LinearWeights vae2llm; + engine::modules::LinearWeights time0; + engine::modules::LinearWeights time2; + core::TensorValue latent_pos_embed; + engine::modules::QwenDecoderStackWeights nar_stack; + engine::modules::NormWeights final_norm; + engine::modules::LinearWeights llm2vae; +}; + +engine::modules::QwenDecoderLayerWeights load_nar_layer( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const Yue2ModelConfig & config, + assets::TensorStorageType storage_type, + int64_t layer) { + const std::string prefix = "model.layers." + std::to_string(layer); + engine::modules::QwenDecoderLayerWeights out; + out.input_norm = binding::norm_weight_from_source(store, source, prefix + ".nar_input_layernorm", config.hidden_size); + out.self_attention.q_weight = store.load_tensor( + source, + prefix + ".nar_self_attn.q_proj.weight", + storage_type, + {config.attention_heads * config.head_dim, config.hidden_size}); + out.self_attention.k_weight = store.load_tensor( + source, + prefix + ".nar_self_attn.k_proj.weight", + storage_type, + {config.kv_heads * config.head_dim, config.hidden_size}); + out.self_attention.v_weight = store.load_tensor( + source, + prefix + ".nar_self_attn.v_proj.weight", + storage_type, + {config.kv_heads * config.head_dim, config.hidden_size}); + out.self_attention.out_weight = store.load_tensor( + source, + prefix + ".nar_self_attn.o_proj.weight", + storage_type, + {config.hidden_size, config.attention_heads * config.head_dim}); + out.q_norm = binding::norm_weight_from_source(store, source, prefix + ".nar_self_attn.q_norm", config.head_dim); + out.k_norm = binding::norm_weight_from_source(store, source, prefix + ".nar_self_attn.k_norm", config.head_dim); + out.post_norm = binding::norm_weight_from_source(store, source, prefix + ".nar_pre_mlp_layernorm", config.hidden_size); + out.mlp.gate_proj = binding::linear_from_source( + store, + source, + prefix + ".nar_mlp.gate_proj", + storage_type, + config.intermediate_size, + config.hidden_size, + false); + out.mlp.up_proj = binding::linear_from_source( + store, + source, + prefix + ".nar_mlp.up_proj", + storage_type, + config.intermediate_size, + config.hidden_size, + false); + out.mlp.down_proj = binding::linear_from_source( + store, + source, + prefix + ".nar_mlp.down_proj", + storage_type, + config.hidden_size, + config.intermediate_size, + false); + return out; +} + +std::shared_ptr load_nar_weights( + const Yue2Assets & assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) { + const auto total_start = Clock::now(); + auto weights = std::make_shared(); + weights->store = std::make_shared( + execution.backend(), + execution.backend_type(), + "yue2.nar.weights", + weight_context_bytes); + const auto & source = *assets.model_weights; + const auto & config = assets.config.model; + const auto bind_start = Clock::now(); + weights->vae2llm = binding::linear_from_source( + *weights->store, + source, + "vae2llm", + storage_type, + config.hidden_size, + config.latent_dim, + true); + weights->time0 = binding::linear_from_source( + *weights->store, + source, + "time_embedder.mlp.0", + storage_type, + config.hidden_size, + 256, + true); + weights->time2 = binding::linear_from_source( + *weights->store, + source, + "time_embedder.mlp.2", + storage_type, + config.hidden_size, + config.hidden_size, + true); + weights->latent_pos_embed = weights->store->load_tensor( + source, + "latent_pos_embed.pe", + storage_type, + {config.max_latent_frames, config.hidden_size}); + weights->nar_stack.layers.reserve(static_cast(config.layers)); + for (int64_t layer = 0; layer < config.layers; ++layer) { + weights->nar_stack.layers.push_back(load_nar_layer(*weights->store, source, config, storage_type, layer)); + } + weights->final_norm = binding::norm_weight_from_source(*weights->store, source, "model.norm", config.hidden_size); + weights->llm2vae = binding::linear_from_source( + *weights->store, + source, + "llm2vae", + storage_type, + config.latent_dim, + config.hidden_size, + true); + engine::debug::timing_log_scalar("yue2.nar.weights_bind_ms", engine::debug::elapsed_ms(bind_start)); + const auto upload_start = Clock::now(); + weights->store->upload(); + engine::debug::timing_log_scalar("yue2.nar.weights_upload_ms", engine::debug::elapsed_ms(upload_start)); + engine::debug::timing_log_scalar("yue2.nar.weights_total_ms", engine::debug::elapsed_ms(total_start)); + return weights; +} + +std::vector> chunk_ranges(int64_t frames, int64_t prefix_tokens, int64_t context) { + const int64_t chunk = std::min((context - prefix_tokens - 3) / 2, context); + if (frames < 1 || chunk < 1) { + throw std::runtime_error("Yue2 NAR chunk leaves no acoustic context"); + } + std::vector> out; + for (int64_t begin = 0; begin < frames; begin += chunk) { + out.push_back({begin, std::min(begin + chunk, frames)}); + } + return out; +} + +std::vector timestep_features(float shifted_t, int64_t rows) { + constexpr int64_t kDim = 256; + std::vector out(static_cast(rows * kDim), 0.0F); + const int64_t half = kDim / 2; + for (int64_t i = 0; i < half; ++i) { + const float freq = std::exp(-std::log(10000.0F) * static_cast(i) / static_cast(half)); + const float value = shifted_t * freq; + for (int64_t row = 0; row < rows; ++row) { + out[static_cast(row * kDim + i)] = std::cos(value); + out[static_cast(row * kDim + half + i)] = std::sin(value); + } + } + return out; +} + +float shifted_t_value(float raw_t, float shift) { + const float sigmoid = 1.0F / (1.0F + std::exp(-raw_t)); + return shift * sigmoid / (1.0F + (shift - 1.0F) * sigmoid); +} + +float logit_clamped(float t) { + const float safe = std::min(std::max(t, 1.0e-7F), 1.0F - 1.0e-7F); + return std::min(std::max(std::log(safe / (1.0F - safe)), -20.0F), 20.0F); +} + +struct QKVParts { + core::TensorValue q; + core::TensorValue k; + core::TensorValue v; +}; + +core::TensorValue reshape_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t heads, + int64_t dim) { + const auto contiguous = core::ensure_backend_addressable_layout(ctx, input); + return core::reshape_tensor( + ctx, + contiguous, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim})); +} + +QKVParts build_qkv_part( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const engine::modules::QwenDecoderLayerWeights & weights, + const Yue2ModelConfig & config) { + auto q = engine::modules::LinearModule({config.hidden_size, config.attention_heads * config.head_dim, false}) + .build(ctx, input, {weights.self_attention.q_weight, std::nullopt}); + auto k = engine::modules::LinearModule({config.hidden_size, config.kv_heads * config.head_dim, false}) + .build(ctx, input, {weights.self_attention.k_weight, std::nullopt}); + auto v = engine::modules::LinearModule({config.hidden_size, config.kv_heads * config.head_dim, false}) + .build(ctx, input, {weights.self_attention.v_weight, std::nullopt}); + q = reshape_heads(ctx, q, config.attention_heads, config.head_dim); + k = reshape_heads(ctx, k, config.kv_heads, config.head_dim); + q = engine::modules::RMSNormModule({config.head_dim, config.rms_norm_eps, true, false}).build(ctx, q, weights.q_norm); + k = engine::modules::RMSNormModule({config.head_dim, config.rms_norm_eps, true, false}).build(ctx, k, weights.k_norm); + return {q, k, reshape_heads(ctx, v, config.kv_heads, config.head_dim)}; +} + +core::TensorValue repeat_kv_heads(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t repeats) { + if (repeats == 1) { + return input; + } + auto contiguous = core::ensure_backend_addressable_layout(ctx, input); + const int64_t batch = contiguous.shape.dims[0]; + const int64_t kv_heads = contiguous.shape.dims[1]; + const int64_t steps = contiguous.shape.dims[2]; + const int64_t dim = contiguous.shape.dims[3]; + auto expanded = core::reshape_tensor( + ctx, + contiguous, + core::TensorShape::from_dims({batch, kv_heads, 1, steps * dim})); + expanded = engine::modules::RepeatModule({core::TensorShape::from_dims({batch, kv_heads, repeats, steps * dim})}) + .build(ctx, expanded); + expanded = core::ensure_backend_addressable_layout(ctx, expanded); + return core::reshape_tensor( + ctx, + expanded, + core::TensorShape::from_dims({batch, kv_heads * repeats, steps, dim})); +} + +core::TensorValue mixed_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & q, + const core::TensorValue & k, + const core::TensorValue & v, + const core::TensorValue * attention_mask, + const Yue2ModelConfig & config, + core::BackendType backend_type) { + auto q_heads = engine::modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + q_heads = core::wrap_tensor(ggml_cont(ctx.ggml, q_heads.tensor), q_heads.shape, q_heads.type); + auto k_heads = engine::modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k); + auto v_heads = engine::modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v); + if (backend_type != core::BackendType::Cpu) { + auto * flash = ggml_flash_attn_ext( + ctx.ggml, + q_heads.tensor, + k_heads.tensor, + v_heads.tensor, + attention_mask != nullptr ? attention_mask->tensor : nullptr, + 1.0F / std::sqrt(static_cast(config.head_dim)), + 0.0F, + 0.0F); + ggml_flash_attn_ext_set_prec(flash, GGML_PREC_F32); + return core::wrap_tensor( + flash, + core::TensorShape::from_dims({q_heads.shape.dims[0], q_heads.shape.dims[2], q_heads.shape.dims[1], config.head_dim}), + GGML_TYPE_F32); + } + + const int64_t repeats = config.attention_heads / config.kv_heads; + k_heads = repeat_kv_heads(ctx, k_heads, repeats); + v_heads = repeat_kv_heads(ctx, v_heads, repeats); + auto scores = engine::modules::MatMulModule{}.build( + ctx, + q_heads, + engine::modules::TransposeModule({{0, 1, 3, 2}, k_heads.shape.rank}).build(ctx, k_heads)); + scores = core::ensure_backend_addressable_layout(ctx, scores); + auto attn = core::wrap_tensor( + ggml_soft_max_ext( + ctx.ggml, + scores.tensor, + attention_mask != nullptr ? attention_mask->tensor : nullptr, + 1.0F / std::sqrt(static_cast(config.head_dim)), + 0.0F), + scores.shape, + GGML_TYPE_F32); + auto context = engine::modules::MatMulModule{}.build(ctx, attn, v_heads); + return engine::modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context); +} + +core::TensorValue build_mlp_part( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const engine::modules::QwenMLPWeights & weights, + const Yue2ModelConfig & config) { + auto gate = engine::modules::LinearModule({config.hidden_size, config.intermediate_size, false}) + .build(ctx, input, {weights.gate_proj.weight, std::nullopt}); + gate = engine::modules::SiluModule{}.build(ctx, gate); + auto up = engine::modules::LinearModule({config.hidden_size, config.intermediate_size, false}) + .build(ctx, input, {weights.up_proj.weight, std::nullopt}); + auto gated = engine::modules::MulModule{}.build(ctx, gate, up); + return engine::modules::LinearModule({config.intermediate_size, config.hidden_size, false}) + .build(ctx, gated, {weights.down_proj.weight, std::nullopt}); +} + +core::TensorValue build_cached_nar_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const engine::modules::QwenDecoderLayerWeights & nar_weights, + const core::TensorValue & ar_key, + const core::TensorValue & ar_value, + const Yue2ModelConfig & config, + core::BackendType backend_type) { + auto norm = engine::modules::RMSNormModule({config.hidden_size, config.rms_norm_eps, true, false}) + .build(ctx, input, nar_weights.input_norm); + auto qkv = build_qkv_part(ctx, norm, nar_weights, config); + qkv.q = engine::modules::RoPEModule({config.head_dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}) + .build(ctx, qkv.q, positions); + qkv.k = engine::modules::RoPEModule({config.head_dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}) + .build(ctx, qkv.k, positions); + qkv.k = core::ensure_backend_addressable_layout(ctx, qkv.k); + qkv.v = core::ensure_backend_addressable_layout(ctx, qkv.v); + qkv.k = core::wrap_tensor(ggml_cast(ctx.ggml, qkv.k.tensor, GGML_TYPE_F16), qkv.k.shape, GGML_TYPE_F16); + qkv.v = core::wrap_tensor(ggml_cast(ctx.ggml, qkv.v.tensor, GGML_TYPE_F16), qkv.v.shape, GGML_TYPE_F16); + auto k = engine::modules::ConcatModule({1}).build(ctx, ar_key, qkv.k); + auto v = engine::modules::ConcatModule({1}).build(ctx, ar_value, qkv.v); + auto context = mixed_attention(ctx, qkv.q, k, v, nullptr, config, backend_type); + context = core::ensure_backend_addressable_layout(ctx, context); + context = core::reshape_tensor( + ctx, + context, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], config.attention_heads * config.head_dim})); + auto attn = engine::modules::LinearModule({config.attention_heads * config.head_dim, config.hidden_size, false}) + .build(ctx, context, {nar_weights.self_attention.out_weight, std::nullopt}); + auto x = engine::modules::AddModule{}.build(ctx, input, attn); + norm = engine::modules::RMSNormModule({config.hidden_size, config.rms_norm_eps, true, false}) + .build(ctx, x, nar_weights.post_norm); + return engine::modules::AddModule{}.build(ctx, x, build_mlp_part(ctx, norm, nar_weights.mlp, config)); +} + +} // namespace + +struct Yue2NarRuntime::Impl { + Impl( + core::ExecutionContext & execution, + std::shared_ptr assets, + assets::TensorStorageType weight_type, + size_t weight_context_bytes, + size_t graph_arena_bytes) + : execution(execution), + assets(std::move(assets)), + graph_arena_bytes(graph_arena_bytes) { + if (!this->assets) { + throw std::runtime_error("Yue2 NAR runtime requires assets"); + } + const auto start = Clock::now(); + weights = load_nar_weights(*this->assets, execution, weight_context_bytes, weight_type); + engine::debug::timing_log_scalar("yue2.nar.init_total_ms", engine::debug::elapsed_ms(start)); + } + + struct HostVelocityTiming { + double pad_ms = 0.0; + double timestep_ms = 0.0; + }; + + struct Graph { + Graph( + Impl & owner, + int64_t frames, + const Yue2ArDevicePrefixState & ar_state) + : owner(&owner), + frames(frames), + ar_length(ar_state.current_end) { + const auto & config = owner.assets->config.model; + if (frames <= 0 || ar_length <= 0) { + throw std::runtime_error("Yue2 NAR graph requires positive shapes"); + } + if (static_cast(ar_state.keys.size()) != config.layers || + static_cast(ar_state.values.size()) != config.layers) { + throw std::runtime_error("Yue2 NAR prefix cache layer count mismatch"); + } + const auto total_start = Clock::now(); + nar_length = frames + 2; + total_length = ar_length + nar_length; + ggml_init_params input_params{ + ggml_tensor_overhead() * static_cast(3 + config.layers * 2), + nullptr, + true}; + input_ctx.reset(ggml_init(input_params)); + ggml_init_params graph_params{owner.graph_arena_bytes, nullptr, true}; + ctx.reset(ggml_init(graph_params)); + if (!input_ctx || !ctx) { + throw std::runtime_error("failed to initialize Yue2 NAR graph context"); + } + core::ModuleBuildContext input_build{input_ctx.get(), "yue2.nar.input", owner.execution.backend_type()}; + core::ModuleBuildContext build{ctx.get(), "yue2.nar", owner.execution.backend_type()}; + const auto build_start = Clock::now(); + state = core::make_tensor(input_build, GGML_TYPE_F32, core::TensorShape::from_dims({1, nar_length, config.latent_dim})); + time = core::make_tensor(input_build, GGML_TYPE_F32, core::TensorShape::from_dims({1, 256})); + positions = core::make_tensor(input_build, GGML_TYPE_I32, core::TensorShape::from_dims({nar_length})); + ggml_set_input(state.tensor); + ggml_set_input(time.tensor); + ggml_set_input(positions.tensor); + ar_keys.reserve(static_cast(config.layers)); + ar_values.reserve(static_cast(config.layers)); + for (int64_t layer = 0; layer < config.layers; ++layer) { + const auto & key = ar_state.keys[static_cast(layer)]; + const auto & value = ar_state.values[static_cast(layer)]; + if (key.shape.rank != 4 || value.shape.rank != 4 || + key.shape.dims[0] != 1 || value.shape.dims[0] != 1 || + key.shape.dims[1] != ar_length || value.shape.dims[1] != ar_length || + key.shape.dims[2] != config.kv_heads || value.shape.dims[2] != config.kv_heads || + key.shape.dims[3] != config.head_dim || value.shape.dims[3] != config.head_dim || + key.type != GGML_TYPE_F16 || value.type != GGML_TYPE_F16) { + throw std::runtime_error("Yue2 NAR prefix cache tensor shape mismatch"); + } + ar_keys.push_back(key); + ar_values.push_back(value); + } + auto nar_hidden = engine::modules::LinearModule({config.latent_dim, config.hidden_size, true}) + .build(build, state, owner.weights->vae2llm); + auto time_hidden = engine::modules::LinearModule({256, config.hidden_size, true}) + .build(build, time, owner.weights->time0); + time_hidden = engine::modules::SiluModule{}.build(build, time_hidden); + time_hidden = engine::modules::LinearModule({config.hidden_size, config.hidden_size, true}) + .build(build, time_hidden, owner.weights->time2); + time_hidden = core::reshape_tensor(build, time_hidden, core::TensorShape::from_dims({1, 1, config.hidden_size})); + time_hidden = engine::modules::RepeatModule( + {core::TensorShape::from_dims({1, nar_length, config.hidden_size})}) + .build(build, time_hidden); + auto pos = engine::modules::SliceModule({0, 0, nar_length}).build(build, owner.weights->latent_pos_embed); + pos = core::reshape_tensor(build, pos, core::TensorShape::from_dims({1, nar_length, config.hidden_size})); + if (pos.type != GGML_TYPE_F32) { + pos = core::wrap_tensor(ggml_cast(build.ggml, pos.tensor, GGML_TYPE_F32), pos.shape, GGML_TYPE_F32); + } + nar_hidden = engine::modules::AddModule{}.build(build, nar_hidden, time_hidden); + nar_hidden = engine::modules::AddModule{}.build(build, nar_hidden, pos); + auto hidden = nar_hidden; + + for (int64_t layer = 0; layer < config.layers; ++layer) { + hidden = build_cached_nar_layer( + build, + hidden, + positions, + owner.weights->nar_stack.layers[static_cast(layer)], + ar_keys[static_cast(layer)], + ar_values[static_cast(layer)], + config, + owner.execution.backend_type()); + } + hidden = engine::modules::RMSNormModule({config.hidden_size, config.rms_norm_eps, true, false}) + .build(build, hidden, owner.weights->final_norm); + auto latent = engine::modules::LinearModule({config.hidden_size, config.latent_dim, true}) + .build(build, hidden, owner.weights->llm2vae); + auto content = engine::modules::SliceModule({1, 1, frames}).build(build, latent); + output = core::ensure_backend_addressable_layout(build, content); + ggml_set_output(output.tensor); + graph = ggml_new_graph_custom(ctx.get(), 262144, false); + ggml_build_forward_expand(graph, output.tensor); + engine::debug::timing_log_scalar("yue2.nar.graph.build_ms", engine::debug::elapsed_ms(build_start)); + const auto alloc_start = Clock::now(); + input_buffer = ggml_backend_alloc_ctx_tensors(input_ctx.get(), owner.execution.backend()); + gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(owner.execution.backend())); + if (input_buffer == nullptr || gallocr == nullptr || + !ggml_gallocr_reserve(gallocr, graph) || + !ggml_gallocr_alloc_graph(gallocr, graph)) { + throw std::runtime_error("failed to allocate Yue2 NAR graph"); + } + engine::debug::timing_log_scalar("yue2.nar.graph.alloc_ms", engine::debug::elapsed_ms(alloc_start)); + std::vector pos_values(static_cast(nar_length)); + for (int64_t i = 0; i < nar_length; ++i) { + pos_values[static_cast(i)] = static_cast(ar_length + i); + } + ggml_backend_tensor_set(positions.tensor, pos_values.data(), 0, pos_values.size() * sizeof(int32_t)); + engine::debug::timing_log_scalar("yue2.nar.graph.static_upload_ms", 0.0); + engine::debug::timing_log_scalar("yue2.nar.graph.frames", frames); + engine::debug::timing_log_scalar("yue2.nar.graph.ar_tokens", ar_length); + engine::debug::timing_log_scalar("yue2.nar.graph.nar_tokens", nar_length); + engine::debug::timing_log_scalar("yue2.nar.graph.total_ms", engine::debug::elapsed_ms(total_start)); + } + + ~Graph() { + core::release_backend_graph_resources(owner->execution.backend(), graph); + if (gallocr) { + ggml_gallocr_free(gallocr); + } + if (input_buffer) { + ggml_backend_buffer_free(input_buffer); + } + } + + bool matches(int64_t f, int64_t ar) const noexcept { + return frames == f && ar_length == ar; + } + + std::vector run( + const std::vector & padded_state, + const std::vector & time_features) { + const auto & config = owner->assets->config.model; + if (static_cast(padded_state.size()) != nar_length * config.latent_dim || + static_cast(time_features.size()) != 256) { + throw std::runtime_error("Yue2 NAR step input shape mismatch"); + } + const auto upload_start = Clock::now(); + core::write_tensor_f32(state, padded_state); + core::write_tensor_f32(time, time_features); + input_upload_ms += engine::debug::elapsed_ms(upload_start); + const auto compute_start = Clock::now(); + const auto status = core::compute_backend_graph(owner->execution.backend(), graph, nullptr, "yue2.nar.velocity"); + ggml_backend_synchronize(owner->execution.backend()); + graph_compute_ms += engine::debug::elapsed_ms(compute_start); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Yue2 NAR graph compute failed"); + } + std::vector out; + const auto read_start = Clock::now(); + core::read_tensor_f32_into(output.tensor, out); + output_read_ms += engine::debug::elapsed_ms(read_start); + ++runs; + return out; + } + + Impl * owner = nullptr; + int64_t frames = 0; + int64_t ar_length = 0; + int64_t nar_length = 0; + int64_t total_length = 0; + std::unique_ptr input_ctx; + std::unique_ptr ctx; + ggml_cgraph * graph = nullptr; + ggml_backend_buffer_t input_buffer = nullptr; + ggml_gallocr_t gallocr = nullptr; + core::TensorValue state; + core::TensorValue time; + core::TensorValue positions; + std::vector ar_keys; + std::vector ar_values; + core::TensorValue output; + double input_upload_ms = 0.0; + double graph_compute_ms = 0.0; + double output_read_ms = 0.0; + int64_t runs = 0; + }; + + std::vector velocity( + Graph & graph, + const Yue2ArDevicePrefixState & ar_state, + const std::vector & state, + float raw_t, + HostVelocityTiming & host_timing) { + const auto & config = assets->config.model; + const int64_t frames = static_cast(state.size()) / config.latent_dim; + if (frames <= 0 || frames * config.latent_dim != static_cast(state.size())) { + throw std::runtime_error("Yue2 NAR state shape mismatch"); + } + if (!graph.matches(frames, ar_state.current_end)) { + throw std::runtime_error("Yue2 NAR graph shape changed during chunk solve"); + } + const auto pad_start = Clock::now(); + std::vector padded(static_cast((frames + 2) * config.latent_dim), 0.0F); + std::copy(state.begin(), state.end(), padded.begin() + static_cast(config.latent_dim)); + host_timing.pad_ms += engine::debug::elapsed_ms(pad_start); + const auto shifted = shifted_t_value(raw_t, config.timestep_shift); + const auto timestep_start = Clock::now(); + auto features = timestep_features(shifted, 1); + host_timing.timestep_ms += engine::debug::elapsed_ms(timestep_start); + return graph.run(padded, features); + } + + std::vector solve_chunk( + const Yue2ArDevicePrefixState & ar_state, + const std::vector & noise, + int64_t ode_steps) { + auto state = noise; + const auto total_start = Clock::now(); + const auto & config = assets->config.model; + const int64_t frames = static_cast(state.size()) / config.latent_dim; + graph = std::make_unique(*this, frames, ar_state); + auto & chunk_graph = *graph; + const float dt = 1.0F / static_cast(ode_steps); + HostVelocityTiming host_timing; + double host_update_ms = 0.0; + for (int64_t step = 0; step < ode_steps; ++step) { + const float t = 1.0F - static_cast(step) * dt; + const auto first = velocity(chunk_graph, ar_state, state, logit_clamped(t), host_timing); + const auto mid_start = Clock::now(); + std::vector mid(state.size(), 0.0F); + for (size_t i = 0; i < state.size(); ++i) { + mid[i] = state[i] - first[i] * (dt / 2.0F); + } + host_update_ms += engine::debug::elapsed_ms(mid_start); + const auto second = velocity(chunk_graph, ar_state, mid, logit_clamped(t - dt / 2.0F), host_timing); + const auto update_start = Clock::now(); + for (size_t i = 0; i < state.size(); ++i) { + state[i] -= second[i] * dt; + } + host_update_ms += engine::debug::elapsed_ms(update_start); + } + engine::debug::timing_log_scalar("yue2.nar.chunk.frames", frames); + engine::debug::timing_log_scalar("yue2.nar.chunk.ode_steps", ode_steps); + engine::debug::timing_log_scalar("yue2.nar.chunk.velocity_runs", chunk_graph.runs); + engine::debug::timing_log_scalar("yue2.nar.chunk.host_pad_ms", host_timing.pad_ms); + engine::debug::timing_log_scalar("yue2.nar.chunk.host_timestep_ms", host_timing.timestep_ms); + engine::debug::timing_log_scalar("yue2.nar.chunk.host_update_ms", host_update_ms); + engine::debug::timing_log_scalar("yue2.nar.chunk.input_upload_ms", chunk_graph.input_upload_ms); + engine::debug::timing_log_scalar("yue2.nar.chunk.graph_compute_ms", chunk_graph.graph_compute_ms); + engine::debug::timing_log_scalar("yue2.nar.chunk.output_read_ms", chunk_graph.output_read_ms); + engine::debug::timing_log_scalar("yue2.nar.chunk.total_ms", engine::debug::elapsed_ms(total_start)); + return state; + } + + std::vector synthesize( + const std::vector & prefix, + const std::vector & codec, + const std::function &)> & prefill_state, + const std::vector & noise, + uint64_t seed, + int64_t ode_steps, + int64_t context) { + const auto total_start = Clock::now(); + const auto & config = assets->config.model; + const auto ranges = chunk_ranges(static_cast(codec.size()), static_cast(prefix.size()), context); + engine::debug::timing_log_scalar("yue2.nar.synthesize.prefix_tokens", prefix.size()); + engine::debug::timing_log_scalar("yue2.nar.synthesize.codec_tokens", codec.size()); + engine::debug::timing_log_scalar("yue2.nar.synthesize.chunks", ranges.size()); + engine::debug::timing_log_scalar("yue2.nar.synthesize.ode_steps", ode_steps); + engine::debug::timing_log_scalar("yue2.nar.synthesize.context", context); + std::vector full_noise(static_cast(codec.size() * config.latent_dim), 0.0F); + const auto noise_start = Clock::now(); + if (noise.empty()) { + std::mt19937 rng(static_cast(seed)); + std::normal_distribution normal(0.0F, 1.0F); + for (float & value : full_noise) { + value = normal(rng); + } + } else { + if (noise.size() != full_noise.size()) { + throw std::runtime_error("Yue2 nar_noise_file frame count does not match semantic codec count"); + } + full_noise = noise; + } + engine::debug::timing_log_scalar("yue2.nar.synthesize.noise_ms", engine::debug::elapsed_ms(noise_start)); + std::vector out; + out.reserve(full_noise.size()); + double token_build_ms = 0.0; + double noise_slice_ms = 0.0; + double prefill_ms = 0.0; + double solve_ms = 0.0; + double append_ms = 0.0; + for (const auto & [begin, end] : ranges) { + const auto token_start = Clock::now(); + std::vector ar_tokens = prefix; + ar_tokens.reserve(prefix.size() + static_cast(end - begin) + 1); + for (int64_t i = begin; i < end; ++i) { + ar_tokens.push_back(codec[static_cast(i)] + kCodecOffset); + } + ar_tokens.push_back(kMusicEndToken); + token_build_ms += engine::debug::elapsed_ms(token_start); + const auto slice_start = Clock::now(); + const size_t begin_elem = static_cast(begin * config.latent_dim); + const size_t end_elem = static_cast(end * config.latent_dim); + std::vector noise(end_elem - begin_elem); + std::copy(full_noise.begin() + static_cast(begin_elem), + full_noise.begin() + static_cast(end_elem), + noise.begin()); + noise_slice_ms += engine::debug::elapsed_ms(slice_start); + const auto prefill_start = Clock::now(); + auto ar_state = prefill_state(ar_tokens); + prefill_ms += engine::debug::elapsed_ms(prefill_start); + const auto solve_start = Clock::now(); + auto chunk = solve_chunk(ar_state, noise, ode_steps); + solve_ms += engine::debug::elapsed_ms(solve_start); + const auto append_start = Clock::now(); + out.insert(out.end(), chunk.begin(), chunk.end()); + append_ms += engine::debug::elapsed_ms(append_start); + } + engine::debug::timing_log_scalar("yue2.nar.synthesize.token_build_ms", token_build_ms); + engine::debug::timing_log_scalar("yue2.nar.synthesize.noise_slice_ms", noise_slice_ms); + engine::debug::timing_log_scalar("yue2.nar.synthesize.prefill_state_ms", prefill_ms); + engine::debug::timing_log_scalar("yue2.nar.synthesize.solve_chunks_ms", solve_ms); + engine::debug::timing_log_scalar("yue2.nar.synthesize.output_append_ms", append_ms); + engine::debug::timing_log_scalar("yue2.nar.synthesize.output_latents", out.size()); + engine::debug::timing_log_scalar("yue2.nar.synthesize.total_ms", engine::debug::elapsed_ms(total_start)); + return out; + } + + core::ExecutionContext & execution; + std::shared_ptr assets; + size_t graph_arena_bytes = 0; + std::shared_ptr weights; + std::unique_ptr graph; +}; + +Yue2NarRuntime::Yue2NarRuntime( + core::ExecutionContext & execution, + std::shared_ptr assets, + assets::TensorStorageType weight_type, + size_t weight_context_bytes, + size_t graph_arena_bytes) + : impl_(std::make_unique( + execution, + std::move(assets), + weight_type, + weight_context_bytes, + graph_arena_bytes)) {} + +Yue2NarRuntime::~Yue2NarRuntime() = default; + +std::vector Yue2NarRuntime::synthesize( + const std::vector & prefix, + const std::vector & codec, + const std::function &)> & prefill_state, + const std::vector & noise, + uint64_t seed, + int64_t ode_steps, + int64_t context) { + return impl_->synthesize(prefix, codec, prefill_state, noise, seed, ode_steps, context); +} + +void Yue2NarRuntime::release_runtime_graphs() { + impl_->graph.reset(); +} + +} // namespace engine::models::yue2 diff --git a/src/models/yue2/pipeline.cpp b/src/models/yue2/pipeline.cpp new file mode 100644 index 00000000..713dd62d --- /dev/null +++ b/src/models/yue2/pipeline.cpp @@ -0,0 +1,466 @@ +#include "engine/models/yue2/pipeline.h" + +#include "engine/framework/debug/profiler.h" +#include "engine/models/yue2/ar_runtime.h" +#include "engine/models/yue2/nar_runtime.h" + +#include +#include +#include +#include +#include + +namespace engine::models::yue2 { +namespace { + +using Clock = std::chrono::steady_clock; + +std::string request_text(const Yue2Request & request) { + std::string text; + text += cot_instruction(request.cot); + text += "\n[Tags]\n"; + text += request.style; + text += "\n[Lyrics]\n"; + text += request.lyrics; + text += "\n"; + return text; +} + +std::vector token_prefixes( + const Yue2Request & request, + const Yue2TextTokenizer & tokenizer, + const std::vector & abc_ids) { + std::vector out; + out.push_back(kEodToken); + auto text_ids = tokenizer.encode(request_text(request)); + out.insert(out.end(), text_ids.begin(), text_ids.end()); + out.push_back(kAbcStartToken); + if (request.cot == Yue2CotMode::Off) { + out.push_back(kAbcEndToken); + out.push_back(kMusicStartToken); + return out; + } + out.insert(out.end(), abc_ids.begin(), abc_ids.end()); + if (!request.abc.empty()) { + out.push_back(kAbcEndToken); + out.push_back(kMusicStartToken); + } + return out; +} + +std::vector negative_prefix( + const Yue2Request & request, + const Yue2TextTokenizer & tokenizer, + const std::vector & abc_ids) { + std::vector out; + out.push_back(kEodToken); + const auto text_ids = tokenizer.encode(cot_instruction(request.cot)); + out.insert(out.end(), text_ids.begin(), text_ids.end()); + if (request.cot == Yue2CotMode::Off) { + out.push_back(kMusicStartToken); + return out; + } + out.push_back(kAbcStartToken); + out.insert(out.end(), abc_ids.begin(), abc_ids.end()); + out.push_back(kAbcEndToken); + out.push_back(kMusicStartToken); + return out; +} + +std::vector codec_from_semantic_tokens(const std::vector & tokens) { + std::vector out; + out.reserve(tokens.size()); + for (const int32_t token : tokens) { + if (token >= kCodecOffset && token < kCodecOffset + kCodecSize) { + out.push_back(token - kCodecOffset); + } + } + return out; +} + +Yue2ArSamplingWindow abc_window(const Yue2GenerationConfig & generation) { + return Yue2ArSamplingWindow{ + 0, + kEodToken, + kAbcEndToken, + generation.abc.min_tokens, + generation.abc.max_tokens, + generation.abc, + }; +} + +Yue2ArSamplingWindow semantic_window(const Yue2GenerationConfig & generation) { + return Yue2ArSamplingWindow{ + kCodecOffset, + kCodecOffset + kCodecSize, + kMusicEndToken, + generation.semantic.min_tokens, + generation.semantic.max_tokens, + generation.semantic, + }; +} + +} // namespace + +class Yue2PipelineRuntime::Impl { +public: + Impl( + core::ExecutionContext & execution, + std::shared_ptr assets, + assets::TensorStorageType model_weight_type, + assets::TensorStorageType vae_weight_type, + size_t model_weight_context_bytes, + size_t vae_weight_context_bytes, + size_t ar_prefill_graph_arena_bytes, + size_t ar_decode_graph_arena_bytes, + size_t nar_graph_arena_bytes, + size_t vae_graph_arena_bytes) + : execution(&execution), + assets(std::move(assets)), + tokenizer(this->assets->tiktoken_path), + model_weight_type(model_weight_type), + vae_weight_type(vae_weight_type), + model_weight_context_bytes(model_weight_context_bytes), + vae_weight_context_bytes(vae_weight_context_bytes), + ar_prefill_graph_arena_bytes(ar_prefill_graph_arena_bytes), + ar_decode_graph_arena_bytes(ar_decode_graph_arena_bytes), + nar_graph_arena_bytes(nar_graph_arena_bytes), + vae_graph_arena_bytes(vae_graph_arena_bytes) { + if (!this->assets) { + throw std::runtime_error("Yue2 pipeline requires assets"); + } + (void) this->nar_graph_arena_bytes; + } + + Yue2Plan plan(const Yue2Request & request) { + Yue2Plan out; + out.cot = request.cot; + if (!request.abc.empty()) { + out.abc = request.abc; + out.abc_ids = tokenizer.encode(request.abc); + } + out.prefix = token_prefixes(request, tokenizer, out.abc_ids); + engine::debug::timing_log_scalar("yue2.plan.prefix_tokens", out.prefix.size()); + engine::debug::timing_log_scalar("yue2.plan.abc_tokens", out.abc_ids.size()); + return out; + } + + Yue2SemanticResult generate_semantic(const Yue2Request & request, Yue2Plan plan) { + const auto total_start = Clock::now(); + Yue2SemanticResult out; + out.plan = std::move(plan); + if (!request.semantic_codes.empty()) { + out.tokens.reserve(request.semantic_codes.size()); + for (const int32_t code : request.semantic_codes) { + out.tokens.push_back(code + kCodecOffset); + } + engine::debug::timing_log_scalar("yue2.semantic.input_codes", request.semantic_codes.size()); + engine::debug::timing_log_scalar("yue2.semantic.tokens", out.tokens.size()); + engine::debug::timing_log_scalar("yue2.semantic.total_inner_ms", engine::debug::elapsed_ms(total_start)); + return out; + } + if (request.cot != Yue2CotMode::Off && request.abc.empty()) { + ensure_ar(); + const auto abc_start = Clock::now(); + out.plan.abc_ids = ar->generate(out.plan.prefix, abc_window(request.generation), request.seed); + engine::debug::timing_log_scalar("yue2.semantic.abc_generate_ms", engine::debug::elapsed_ms(abc_start)); + out.plan.truncated = static_cast(out.plan.abc_ids.size()) >= request.generation.abc.max_tokens; + out.plan.prefix.insert(out.plan.prefix.end(), out.plan.abc_ids.begin(), out.plan.abc_ids.end()); + out.plan.prefix.push_back(kAbcEndToken); + out.plan.prefix.push_back(kMusicStartToken); + engine::debug::timing_log_scalar("yue2.semantic.abc_generated_tokens", out.plan.abc_ids.size()); + engine::debug::timing_log_scalar("yue2.semantic.abc_truncated", out.plan.truncated); + ar->release_runtime_graphs(); + } + ensure_ar(); + const auto negative_start = Clock::now(); + const auto neg = negative_prefix(request, tokenizer, out.plan.abc_ids); + engine::debug::timing_log_scalar("yue2.semantic.negative_prefix_ms", engine::debug::elapsed_ms(negative_start)); + engine::debug::timing_log_scalar("yue2.semantic.positive_prefix_tokens", out.plan.prefix.size()); + engine::debug::timing_log_scalar("yue2.semantic.negative_prefix_tokens", neg.size()); + const auto music_start = Clock::now(); + out.tokens = ar->generate_cfg( + out.plan.prefix, + neg, + semantic_window(request.generation), + request_guidance_scale(request), + request.seed); + engine::debug::timing_log_scalar("yue2.semantic.music_generate_ms", engine::debug::elapsed_ms(music_start)); + out.truncated = static_cast(out.tokens.size()) >= request.generation.semantic.max_tokens; + engine::debug::timing_log_scalar("yue2.semantic.tokens", out.tokens.size()); + engine::debug::timing_log_scalar("yue2.semantic.truncated", out.truncated); + engine::debug::timing_log_scalar("yue2.semantic.total_inner_ms", engine::debug::elapsed_ms(total_start)); + ar->release_runtime_graphs(); + return out; + } + + std::vector synthesize_latents( + const Yue2SemanticResult & semantic, + const Yue2GenerationConfig & generation, + const std::vector & noise, + uint64_t seed) { + const auto codec_start = Clock::now(); + const auto codec = codec_from_semantic_tokens(semantic.tokens); + engine::debug::timing_log_scalar("yue2.nar.codec_extract_ms", engine::debug::elapsed_ms(codec_start)); + engine::debug::timing_log_scalar("yue2.nar.codec_tokens", codec.size()); + if (codec.empty()) { + throw std::runtime_error("Yue2 semantic generation produced no codec tokens"); + } + ensure_nar(); + ensure_ar(); + return nar->synthesize( + semantic.plan.prefix, + codec, + [this](const std::vector & tokens) { + return ar->prefill_device_state(tokens); + }, + noise, + seed, + generation.ode_steps, + generation.context); + } + + runtime::AudioBuffer decode_audio(const std::vector & latents, int64_t frames) { + ensure_vae(); + const int64_t channels = assets->config.model.latent_dim; + if (frames <= 0 || static_cast(latents.size()) != frames * channels) { + throw std::runtime_error("Yue2 VAE latent shape mismatch"); + } + const int64_t core_frames = assets->config.vae.decode_core_frames; + const int64_t halo_frames = assets->config.vae.decode_halo_frames; + const int64_t ratio = assets->config.vae.downsampling_ratio; + if (core_frames <= 0 || halo_frames < 0 || ratio <= 0) { + throw std::runtime_error("Yue2 VAE tile configuration is invalid"); + } + engine::debug::timing_log_scalar("yue2.vae_decode.latent_frames", frames); + engine::debug::timing_log_scalar("yue2.vae_decode.channels", channels); + double planar_pack_ms = 0.0; + + auto make_planar_tile = [&](int64_t begin, int64_t end) { + const auto start = Clock::now(); + const int64_t tile_frames = end - begin; + std::vector planar(static_cast(channels * tile_frames), 0.0F); + for (int64_t t = 0; t < tile_frames; ++t) { + for (int64_t c = 0; c < channels; ++c) { + planar[static_cast(c * tile_frames + t)] = + latents[static_cast((begin + t) * channels + c)]; + } + } + planar_pack_ms += engine::debug::elapsed_ms(start); + return planar; + }; + + if (frames <= core_frames) { + const auto decode_start = Clock::now(); + auto audio = vae->decode(make_planar_tile(0, frames), 1, frames).front(); + engine::debug::timing_log_scalar("yue2.vae_decode.tiles", 1); + engine::debug::timing_log_scalar("yue2.vae_decode.planar_pack_ms", planar_pack_ms); + engine::debug::timing_log_scalar("yue2.vae_decode.tile_decode_ms", engine::debug::elapsed_ms(decode_start)); + engine::debug::timing_log_scalar("yue2.vae_decode.tile_copy_ms", 0.0); + engine::debug::timing_log_scalar("yue2.vae_decode.output_frames", static_cast(audio.samples.size()) / audio.channels); + return audio; + } + + const int64_t total_output_frames = frames * ratio - 64; + if (total_output_frames <= 0) { + throw std::runtime_error("Yue2 VAE output frame count is invalid"); + } + runtime::AudioBuffer audio; + audio.sample_rate = assets->config.vae.sample_rate; + audio.channels = static_cast(assets->config.vae.channels); + audio.samples.assign(static_cast(total_output_frames * audio.channels), 0.0F); + int64_t tiles = 0; + double tile_decode_ms = 0.0; + double tile_copy_ms = 0.0; + for (int64_t start = 0; start < frames; start += core_frames) { + const int64_t end = std::min(frames, start + core_frames); + const int64_t left = std::max(0, start - halo_frames); + const int64_t right = std::min(frames, end + halo_frames); + const auto tile_decode_start = Clock::now(); + auto tile_audio = vae->decode(make_planar_tile(left, right), 1, right - left).front(); + tile_decode_ms += engine::debug::elapsed_ms(tile_decode_start); + const int64_t out_start = start * ratio; + const int64_t out_end = std::min(end * ratio, total_output_frames); + const int64_t copy_frames = out_end - out_start; + const int64_t crop_start = (start - left) * ratio; + if (copy_frames <= 0 || + crop_start < 0 || + crop_start + copy_frames > static_cast(tile_audio.samples.size()) / tile_audio.channels) { + throw std::runtime_error("Yue2 VAE tile did not cover output core"); + } + const auto copy_start = Clock::now(); + for (int64_t t = 0; t < copy_frames; ++t) { + for (int64_t c = 0; c < audio.channels; ++c) { + audio.samples[static_cast((out_start + t) * audio.channels + c)] = + tile_audio.samples[static_cast((crop_start + t) * tile_audio.channels + c)]; + } + } + tile_copy_ms += engine::debug::elapsed_ms(copy_start); + ++tiles; + } + engine::debug::timing_log_scalar("yue2.vae_decode.tiles", tiles); + engine::debug::timing_log_scalar("yue2.vae_decode.planar_pack_ms", planar_pack_ms); + engine::debug::timing_log_scalar("yue2.vae_decode.tile_decode_ms", tile_decode_ms); + engine::debug::timing_log_scalar("yue2.vae_decode.tile_copy_ms", tile_copy_ms); + engine::debug::timing_log_scalar("yue2.vae_decode.output_frames", total_output_frames); + return audio; + } + + runtime::AudioBuffer run(const Yue2Request & request) { + const auto plan_start = Clock::now(); + auto planned = plan(request); + engine::debug::timing_log_scalar("yue2.plan_ms", engine::debug::elapsed_ms(plan_start, Clock::now())); + const auto semantic_start = Clock::now(); + auto semantic = generate_semantic(request, std::move(planned)); + engine::debug::timing_log_scalar("yue2.semantic_ms", engine::debug::elapsed_ms(semantic_start, Clock::now())); + const auto nar_start = Clock::now(); + auto latents = synthesize_latents(semantic, request.generation, request.nar_noise, request.seed); + engine::debug::timing_log_scalar("yue2.nar_ms", engine::debug::elapsed_ms(nar_start, Clock::now())); + const int64_t frames = static_cast(latents.size()) / assets->config.model.latent_dim; + ar.reset(); + nar.reset(); + const auto vae_start = Clock::now(); + auto audio = decode_audio(latents, frames); + engine::debug::timing_log_scalar("yue2.vae_decode_ms", engine::debug::elapsed_ms(vae_start, Clock::now())); + if (vae) { + vae->release_runtime_graphs(); + } + return audio; + } + + void release_runtime_graphs() { + if (vae) { + vae->release_runtime_graphs(); + } + if (ar) { + ar->release_runtime_graphs(); + } + if (nar) { + nar->release_runtime_graphs(); + } + } + +private: + void ensure_ar() { + if (ar) { + return; + } + const auto start = Clock::now(); + ar = std::make_unique( + *execution, + assets, + model_weight_type, + model_weight_context_bytes, + ar_prefill_graph_arena_bytes, + ar_decode_graph_arena_bytes); + engine::debug::timing_log_scalar("yue2.ar.init_ms", engine::debug::elapsed_ms(start)); + } + + void ensure_vae() { + if (vae) { + return; + } + codecs::OobleckAudioVaeConfig config; + config.sample_rate = assets->config.vae.sample_rate; + config.audio_channels = assets->config.vae.channels; + config.encoder_latent_dim = assets->config.vae.encoder_latent_dim; + config.decoder_latent_dim = assets->config.vae.latent_dim; + config.encoder_prefix = "encoder"; + config.decoder_prefix = "decoder"; + codecs::OobleckAudioVaeRuntimeOptions options; + options.weight_context_bytes = vae_weight_context_bytes; + options.graph_arena_bytes = vae_graph_arena_bytes; + options.weight_storage_type = vae_weight_type; + const auto start = Clock::now(); + vae = std::make_unique( + assets->vae_weights, + *execution, + std::move(config), + options); + engine::debug::timing_log_scalar("yue2.vae.init_ms", engine::debug::elapsed_ms(start)); + } + + void ensure_nar() { + if (nar) { + return; + } + const auto start = Clock::now(); + nar = std::make_unique( + *execution, + assets, + model_weight_type, + model_weight_context_bytes, + nar_graph_arena_bytes); + engine::debug::timing_log_scalar("yue2.nar.init_ms", engine::debug::elapsed_ms(start)); + } + + core::ExecutionContext * execution = nullptr; + std::shared_ptr assets; + Yue2TextTokenizer tokenizer; + assets::TensorStorageType model_weight_type = assets::TensorStorageType::Native; + assets::TensorStorageType vae_weight_type = assets::TensorStorageType::Native; + size_t model_weight_context_bytes = 0; + size_t vae_weight_context_bytes = 0; + size_t ar_prefill_graph_arena_bytes = 0; + size_t ar_decode_graph_arena_bytes = 0; + size_t nar_graph_arena_bytes = 0; + size_t vae_graph_arena_bytes = 0; + std::unique_ptr vae; + std::unique_ptr ar; + std::unique_ptr nar; +}; + +Yue2PipelineRuntime::Yue2PipelineRuntime( + core::ExecutionContext & execution, + std::shared_ptr assets, + assets::TensorStorageType model_weight_type, + assets::TensorStorageType vae_weight_type, + size_t model_weight_context_bytes, + size_t vae_weight_context_bytes, + size_t ar_prefill_graph_arena_bytes, + size_t ar_decode_graph_arena_bytes, + size_t nar_graph_arena_bytes, + size_t vae_graph_arena_bytes) + : impl_(std::make_unique( + execution, + std::move(assets), + model_weight_type, + vae_weight_type, + model_weight_context_bytes, + vae_weight_context_bytes, + ar_prefill_graph_arena_bytes, + ar_decode_graph_arena_bytes, + nar_graph_arena_bytes, + vae_graph_arena_bytes)) {} + +Yue2PipelineRuntime::~Yue2PipelineRuntime() = default; + +Yue2Plan Yue2PipelineRuntime::plan(const Yue2Request & request) { + return impl_->plan(request); +} + +Yue2SemanticResult Yue2PipelineRuntime::generate_semantic(const Yue2Request & request, Yue2Plan plan) { + return impl_->generate_semantic(request, std::move(plan)); +} + +std::vector Yue2PipelineRuntime::synthesize_latents( + const Yue2SemanticResult & semantic, + const Yue2GenerationConfig & generation, + uint64_t seed) { + static const std::vector empty_noise; + return impl_->synthesize_latents(semantic, generation, empty_noise, seed); +} + +runtime::AudioBuffer Yue2PipelineRuntime::decode_audio(const std::vector & latents, int64_t frames) { + return impl_->decode_audio(latents, frames); +} + +runtime::AudioBuffer Yue2PipelineRuntime::run(const Yue2Request & request) { + return impl_->run(request); +} + +void Yue2PipelineRuntime::release_runtime_graphs() { + impl_->release_runtime_graphs(); +} + +} // namespace engine::models::yue2 diff --git a/src/models/yue2/request.cpp b/src/models/yue2/request.cpp new file mode 100644 index 00000000..e6e547e4 --- /dev/null +++ b/src/models/yue2/request.cpp @@ -0,0 +1,192 @@ +#include "engine/models/yue2/request.h" + +#include "engine/framework/io/filesystem.h" +#include "engine/framework/io/binary.h" +#include "engine/framework/runtime/options.h" + +#include +#include +#include +#include + +namespace engine::models::yue2 { +namespace { + +std::string request_style(const runtime::TaskRequest & request) { + if (const auto style = runtime::find_option(request.options, {"style"})) { + return *style; + } + if (request.voice.has_value() && request.voice->style.has_value()) { + const auto & tags = request.voice->style->tags; + if (const auto it = tags.find("style"); it != tags.end()) { + return it->second; + } + if (const auto it = tags.find("tags"); it != tags.end()) { + return it->second; + } + } + return {}; +} + +std::string request_lyrics(const runtime::TaskRequest & request) { + if (const auto lyrics = runtime::find_option(request.options, {"lyrics"})) { + return *lyrics; + } + if (request.text_input.has_value()) { + return request.text_input->text; + } + return {}; +} + +std::string abc_from_options(const std::unordered_map & options) { + if (const auto abc = runtime::find_option(options, {"abc"})) { + return *abc; + } + if (const auto abc_file = runtime::find_option(options, {"abc_file"})) { + const std::filesystem::path path(*abc_file); + if (!engine::io::is_existing_file(path)) { + throw std::runtime_error("Yue2 abc_file does not exist: " + path.string()); + } + return engine::io::read_text_file(path); + } + return {}; +} + +void apply_options( + Yue2Request & out, + const std::unordered_map & options) { + if (const auto cot = runtime::find_option(options, {"cot"})) { + out.cot = parse_cot_mode(*cot); + } + if (const auto seed = runtime::parse_u64_option(options, {"seed"})) { + out.seed = *seed; + } + if (out.seed >= (uint64_t{1} << 63U)) { + throw std::runtime_error("Yue2 seed must be in [0, 2^63)"); + } + if (const auto value = runtime::parse_finite_float_option(options, {"cfg_scale"})) { + if (*value < 0.0F || *value > 20.0F) { + throw std::runtime_error("Yue2 cfg_scale must be in [0,20]"); + } + out.cfg_scale = *value; + } + out.generation.ode_steps = + runtime::parse_positive_i64_option(options, {"num_inference_steps"}, out.generation.ode_steps); + if (out.generation.ode_steps <= 0) { + throw std::runtime_error("Yue2 num_inference_steps must be positive"); + } + auto apply_sampling = [&options](Yue2SamplingConfig & sampling, const std::string & prefix) { + if (const auto value = runtime::parse_finite_float_option(options, {prefix + "_temperature"})) { + sampling.temperature = *value; + } + if (const auto value = runtime::parse_finite_float_option(options, {prefix + "_top_p"})) { + sampling.top_p = *value; + } + if (const auto value = runtime::parse_i64_option(options, {prefix + "_top_k"})) { + sampling.top_k = *value; + } + if (const auto value = runtime::parse_finite_float_option(options, {prefix + "_repetition_penalty"})) { + sampling.repetition_penalty = *value; + } + if (const auto value = runtime::parse_i64_option(options, {prefix + "_penalty_window"})) { + sampling.penalty_window = *value; + } + if (const auto value = runtime::parse_i64_option(options, {prefix + "_min_tokens"})) { + sampling.min_tokens = *value; + } + if (const auto value = runtime::parse_i64_option(options, {prefix + "_max_tokens"})) { + sampling.max_tokens = *value; + } + }; + auto validate_sampling = [](const Yue2SamplingConfig & sampling, const std::string & prefix) { + if (sampling.temperature < 0.0F || sampling.temperature > 5.0F || + sampling.top_p <= 0.0F || sampling.top_p > 1.0F || + sampling.top_k < 1 || + sampling.repetition_penalty <= 0.0F || + sampling.penalty_window < 1 || + sampling.min_tokens < 0 || + sampling.max_tokens < sampling.min_tokens) { + throw std::runtime_error("Yue2 " + prefix + " sampling options are invalid"); + } + }; + apply_sampling(out.generation.abc, "abc"); + apply_sampling(out.generation.semantic, "semantic"); + validate_sampling(out.generation.abc, "abc"); + validate_sampling(out.generation.semantic, "semantic"); + out.abc = abc_from_options(options); + if (!out.abc.empty() && out.cot == Yue2CotMode::Off) { + throw std::runtime_error("Yue2 external ABC requires cot=melody or cot=full"); + } + if (const auto semantic_codes_file = runtime::find_option(options, {"semantic_codes_file"})) { + const std::filesystem::path path(*semantic_codes_file); + if (!engine::io::is_existing_file(path)) { + throw std::runtime_error("Yue2 semantic_codes_file does not exist: " + path.string()); + } + out.semantic_codes = engine::io::read_i32_file(path); + if (out.semantic_codes.empty()) { + throw std::runtime_error("Yue2 semantic_codes_file is empty: " + path.string()); + } + for (const int32_t code : out.semantic_codes) { + if (code < 0 || code >= kCodecSize) { + throw std::runtime_error("Yue2 semantic_codes_file contains a code outside [0,32768)"); + } + } + } + if (const auto nar_noise_file = runtime::find_option(options, {"nar_noise_file"})) { + const std::filesystem::path path(*nar_noise_file); + if (!engine::io::is_existing_file(path)) { + throw std::runtime_error("Yue2 nar_noise_file does not exist: " + path.string()); + } + out.nar_noise = engine::io::read_f32_file(path); + if (out.nar_noise.empty()) { + throw std::runtime_error("Yue2 nar_noise_file is empty: " + path.string()); + } + if (out.nar_noise.size() % static_cast(Yue2ModelConfig{}.latent_dim) != 0) { + throw std::runtime_error("Yue2 nar_noise_file must contain raw float32 acoustic noise rows with 64 columns"); + } + } +} + +Yue2Request normalize_request(Yue2Request out) { + if (out.style.empty()) { + throw std::runtime_error("Yue2 requires non-empty style"); + } + if (out.lyrics.empty()) { + throw std::runtime_error("Yue2 requires non-empty lyrics"); + } + return out; +} + +} // namespace + +Yue2Request parse_yue2_request(const runtime::TaskRequest & request, const Yue2GenerationConfig & defaults) { + if (request.audio_input.has_value()) { + throw std::runtime_error("Yue2 does not consume audio_input"); + } + if (!request.input_artifacts.empty()) { + throw std::runtime_error("Yue2 does not consume input artifacts"); + } + Yue2Request out; + out.generation = defaults; + out.style = request_style(request); + out.lyrics = request_lyrics(request); + apply_options(out, request.options); + return normalize_request(std::move(out)); +} + +Yue2Request parse_yue2_preparation_request( + const runtime::SessionPreparationRequest & request, + const Yue2GenerationConfig & defaults) { + Yue2Request out; + out.generation = defaults; + if (request.text.has_value()) { + out.lyrics = request.text->text; + } + if (const auto style = runtime::find_option(request.options, {"style"})) { + out.style = *style; + } + apply_options(out, request.options); + return out; +} + +} // namespace engine::models::yue2 diff --git a/src/models/yue2/session.cpp b/src/models/yue2/session.cpp new file mode 100644 index 00000000..234fe619 --- /dev/null +++ b/src/models/yue2/session.cpp @@ -0,0 +1,309 @@ +#include "engine/models/yue2/session.h" + +#include "engine/framework/debug/profiler.h" +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/io/filesystem.h" +#include "engine/framework/runtime/options.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::yue2 { +namespace { + +using Clock = std::chrono::steady_clock; +constexpr const char * kFamily = "yue2"; + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("Yue2 session requires assets"); + } + return assets; +} + +engine::assets::TensorStorageType parse_weight_type( + const runtime::SessionOptions & options, + const char * key, + assets::TensorStorageType fallback) { + return runtime::parse_tensor_storage_option( + options.options, + key, + "yue2.weight_type", + fallback, + { + assets::TensorStorageType::Native, + assets::TensorStorageType::F32, + assets::TensorStorageType::F16, + assets::TensorStorageType::BF16, + assets::TensorStorageType::Q8_0, + assets::TensorStorageType::Q4_0, + assets::TensorStorageType::Q4_K, + }); +} + +std::filesystem::path resolve_component_gguf_path( + const Yue2Assets & assets, + std::string_view option_name, + const std::string & value) { + if (value.empty()) { + throw std::runtime_error(std::string(option_name) + " must not be empty"); + } + const std::filesystem::path relative(value); + if (relative.is_absolute()) { + throw std::runtime_error(std::string(option_name) + " must be relative to the Yue2 model root"); + } + const auto path = assets.model_root / relative; + if (!engine::io::is_existing_file(path)) { + throw std::runtime_error(std::string(option_name) + " file does not exist: " + path.string()); + } + if (path.extension() != ".gguf") { + throw std::runtime_error(std::string(option_name) + " must point to a GGUF file"); + } + return path; +} + +void validate_component_anchors(const Yue2Assets & assets) { + const auto & source = *assets.model_weights; + assets::require_tensor_shape( + source, + "model.embed_tokens.weight", + {assets.config.model.vocab_size, assets.config.model.hidden_size}); + assets::require_tensor_shape( + source, + "model.layers.0.self_attn.q_proj.weight", + {assets.config.model.attention_heads * assets.config.model.head_dim, assets.config.model.hidden_size}); + assets::require_tensor_shape( + source, + "model.layers.0.nar_self_attn.q_proj.weight", + {assets.config.model.attention_heads * assets.config.model.head_dim, assets.config.model.hidden_size}); + assets::require_tensor_shape(source, "vae2llm.weight", {assets.config.model.hidden_size, assets.config.model.latent_dim}); + assets::require_tensor_shape(source, "llm2vae.weight", {assets.config.model.latent_dim, assets.config.model.hidden_size}); + if (assets.vae_weights->has_tensor("decoder.layers.0.weight")) { + assets::require_tensor_shape(*assets.vae_weights, "decoder.layers.0.weight", {2048, 64, 7}); + } else { + assets::require_tensor_shape(*assets.vae_weights, "decoder.layers.0.weight_g", {2048, 1, 1}); + } +} + +std::shared_ptr select_component_assets( + std::shared_ptr base, + const std::unordered_map & options) { + auto selected = std::make_shared(*base); + const std::string model_gguf = + runtime::find_option(options, {"yue2.model_gguf"}).value_or("yue2-3b-q8_0.gguf"); + selected->model_weights = assets::open_tensor_source( + resolve_component_gguf_path(*base, "yue2.model_gguf", model_gguf), + "model_weights"); + const std::string vae_gguf = + runtime::find_option(options, {"yue2.vae_gguf"}).value_or("yue2-vae-f16.gguf"); + selected->vae_weights = assets::open_tensor_source( + resolve_component_gguf_path(*base, "yue2.vae_gguf", vae_gguf), + "vae_weights"); + validate_component_anchors(*selected); + return selected; +} + +std::unique_ptr create_yue2_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options, + std::shared_ptr assets) { + return std::make_unique( + task, + options, + std::move(assets)); +} + +const runtime::ModelMetadata & yue2_metadata() noexcept { + static const runtime::ModelMetadata metadata{ + kFamily, + "Yue2", + "Yue2 music generation with lyrics, style, and optional ABC conditioning.", + { + "sidecars/yue2-model-config.json", + "sidecars/yue2-generation-config.json", + "sidecars/yue2-qwen.tiktoken", + "sidecars/yue2-vae-config.json", + }, + { + "yue2-3b-q8_0.gguf", + "yue2-3b-q4_0.gguf", + "yue2-3b-bf16.gguf", + "yue2-vae-f16.gguf", + "yue2-vae-f32.gguf", + }}; + return metadata; +} + +const runtime::CapabilitySet & yue2_capabilities() noexcept { + static const runtime::CapabilitySet capabilities{ + {runtime::TaskCapability{ + runtime::VoiceTaskKind::AudioGeneration, + {runtime::RunMode::Offline}, + }}, + {"auto"}, + false, + true, + false, + }; + return capabilities; +} + +runtime::ModelCliInterface yue2_cli_interface() { + runtime::ModelCliInterface out; + out.request_options = { + {"style", "string", "Music style prompt.", true}, + {"lyrics", "string", "Lyrics to generate.", false}, + {"abc", "string", "ABC score conditioning text.", false}, + {"abc_file", "path", "Path to ABC score conditioning text.", false}, + {"cot", "off|melody|full", "Planning mode.", false, "off"}, + {"seed", "int", "Generation seed.", false, "0"}, + {"cfg_scale", "float", "Classifier-free guidance scale.", false, "1.0", "0.0", "20.0"}, + {"num_inference_steps", "int", "NAR ODE steps.", false, "10", "1"}, + }; + out.session_options = { + {"yue2.model_gguf", "string", "Yue2 main AR/NAR component GGUF file relative to the model root.", false, "yue2-3b-q8_0.gguf"}, + {"yue2.vae_gguf", "string", "Yue2 VAE component GGUF file relative to the model root.", false, "yue2-vae-f16.gguf"}, + {"yue2.model_weight_type", "native|f32|f16|bf16|q8_0|q4_0|q4_k", "Yue2 main model weight storage type.", false, "native"}, + {"yue2.vae_weight_type", "native|f32|f16|bf16|q8_0|q4_0|q4_k", "Yue2 VAE weight storage type.", false, "native"}, + {"yue2.model_weight_context_mb", "int", "Yue2 main model weight context size in MiB.", false, "6144", "1"}, + {"yue2.vae_weight_context_mb", "int", "Yue2 VAE weight context size in MiB.", false, "1536", "1"}, + {"yue2.ar_prefill_graph_arena_mb", "int", "AR prefill graph arena size in MiB.", false, "4096", "1"}, + {"yue2.ar_decode_graph_arena_mb", "int", "AR one-token decode graph arena size in MiB.", false, "1536", "1"}, + {"yue2.nar_graph_arena_mb", "int", "NAR acoustic flow graph arena size in MiB.", false, "6144", "1"}, + {"yue2.vae_graph_arena_mb", "int", "VAE decode graph arena size in MiB.", false, "1536", "1"}, + }; + return out; +} + +} // namespace + +Yue2Session::Yue2Session( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets) + : runtime::RuntimeSessionBase(options), + task_(task), + assets_(select_component_assets(require_assets(std::move(assets)), options.options)) { + if (task_.task != runtime::VoiceTaskKind::AudioGeneration || task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Yue2 supports only offline gen/music"); + } + pipeline_ = std::make_unique( + execution_context(), + assets_, + parse_weight_type(options, "yue2.model_weight_type", assets::TensorStorageType::Native), + parse_weight_type(options, "yue2.vae_weight_type", assets::TensorStorageType::Native), + runtime::parse_size_mb_option(options.options, {"yue2.model_weight_context_mb"}, 6144ull * 1024ull * 1024ull), + runtime::parse_size_mb_option(options.options, {"yue2.vae_weight_context_mb"}, 1536ull * 1024ull * 1024ull), + runtime::parse_size_mb_option(options.options, {"yue2.ar_prefill_graph_arena_mb"}, 4096ull * 1024ull * 1024ull), + runtime::parse_size_mb_option(options.options, {"yue2.ar_decode_graph_arena_mb"}, 1536ull * 1024ull * 1024ull), + runtime::parse_size_mb_option(options.options, {"yue2.nar_graph_arena_mb"}, 6144ull * 1024ull * 1024ull), + runtime::parse_size_mb_option(options.options, {"yue2.vae_graph_arena_mb"}, 1536ull * 1024ull * 1024ull)); +} + +Yue2Session::~Yue2Session() = default; + +std::string Yue2Session::family() const { + return kFamily; +} + +runtime::VoiceTaskKind Yue2Session::task_kind() const { + return task_.task; +} + +runtime::RunMode Yue2Session::run_mode() const { + return task_.mode; +} + +void Yue2Session::prepare(const runtime::SessionPreparationRequest & request) { + if (request.text.has_value() || !request.options.empty()) { + (void) parse_yue2_preparation_request(request, assets_->config.generation); + } + mark_prepared(); +} + +runtime::TaskResult Yue2Session::run(const runtime::TaskRequest & request) { + require_prepared("Yue2 run"); + const auto wall_start = Clock::now(); + const auto parsed = parse_yue2_request(request, assets_->config.generation); + runtime::TaskResult result; + result.audio_output = pipeline_->run(parsed); + engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start, Clock::now())); + return result; +} + +std::shared_ptr make_yue2_loader() { + class LoadedModel final : public runtime::ILoadedVoiceModel { + public: + explicit LoadedModel(std::shared_ptr assets) + : assets_(require_assets(std::move(assets))) {} + + const runtime::ModelMetadata & metadata() const noexcept override { + return yue2_metadata(); + } + + const runtime::CapabilitySet & capabilities() const noexcept override { + return yue2_capabilities(); + } + + std::unique_ptr create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const override { + return create_yue2_session(task, options, assets_); + } + + private: + std::shared_ptr assets_; + }; + + class Loader final : public runtime::IVoiceModelLoader { + public: + std::string family() const override { + return kFamily; + } + + bool can_load(const runtime::ModelLoadRequest & request) const override { + if (request.family_hint.has_value() && *request.family_hint != kFamily) { + return false; + } + try { + (void) load_yue2_assets(request.model_path); + return true; + } catch (const std::exception &) { + return false; + } + } + + runtime::ModelInspection inspect(const runtime::ModelLoadRequest & request) const override { + auto assets = load_yue2_assets(request.model_path); + runtime::ModelInspection inspection; + inspection.model_root = assets->model_root; + inspection.metadata = yue2_metadata(); + inspection.capabilities = yue2_capabilities(); + inspection.cli = yue2_cli_interface(); + inspection.discovered_configs = runtime::discover_named_assets( + inspection.model_root, + inspection.metadata.config_candidates); + inspection.discovered_weights = runtime::discover_named_assets( + inspection.model_root, + inspection.metadata.weight_candidates); + return inspection; + } + + std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { + return std::make_unique(load_yue2_assets(request.model_path)); + } + + runtime::CapabilitySet advertised_capabilities() const override { + return yue2_capabilities(); + } + }; + + return std::make_shared(); +} + +} // namespace engine::models::yue2 diff --git a/src/models/yue2/tokenizer_text.cpp b/src/models/yue2/tokenizer_text.cpp new file mode 100644 index 00000000..56d5b965 --- /dev/null +++ b/src/models/yue2/tokenizer_text.cpp @@ -0,0 +1,153 @@ +#include "engine/models/yue2/tokenizer_text.h" + +#include "engine/models/yue2/types.h" + +#include "bpe-core.h" +#include "unicode.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::yue2 { +namespace { + +namespace vendor = llama_tokenizer_vendor; + +std::string decode_base64(const std::string & input) { + static const std::array table = [] { + std::array values{}; + values.fill(-1); + for (int i = 0; i < 26; ++i) { + values[static_cast('A' + i)] = static_cast(i); + values[static_cast('a' + i)] = static_cast(26 + i); + } + for (int i = 0; i < 10; ++i) { + values[static_cast('0' + i)] = static_cast(52 + i); + } + values[static_cast('+')] = 62; + values[static_cast('/')] = 63; + return values; + }(); + std::string out; + int bits = 0; + int value = 0; + for (const unsigned char ch : input) { + if (ch == '=') { + break; + } + const int8_t digit = table[ch]; + if (digit < 0) { + throw std::runtime_error("Yue2 tiktoken vocabulary contains invalid base64 bytes"); + } + value = (value << 6) | digit; + bits += 6; + if (bits >= 8) { + bits -= 8; + out.push_back(static_cast((value >> bits) & 0xff)); + } + } + return out; +} + +std::string map_token_bytes(const std::string & bytes) { + std::string mapped; + for (const unsigned char byte : bytes) { + mapped += unicode_byte_to_utf8(byte); + } + return mapped; +} + +std::string pair_key(const std::string & left, const std::string & right) { + std::string key = left; + key.push_back('\0'); + key += right; + return key; +} + +void add_special_token(vendor::BpeVocabulary & vocab, const std::string & text, int32_t id) { + vocab.token_to_id.emplace(text, id); + vocab.id_to_token.emplace(id, vendor::TokenData{text, vendor::TOKEN_ATTR_CONTROL}); +} + +void register_yue2_special_tokens(vendor::BpeVocabulary & vocab, int32_t base_id) { + static const std::array kBaseSpecials = { + "<|endoftext|>", + "<|im_start|>", + "<|im_end|>", + "", + "", + "", + "", + "", + }; + int32_t id = base_id; + for (const char * token : kBaseSpecials) { + add_special_token(vocab, token, id++); + } + for (int i = 0; i < 200; ++i) { + add_special_token(vocab, "", id++); + } + add_special_token(vocab, "", kAbcStartToken); + add_special_token(vocab, "", kAbcEndToken); + add_special_token(vocab, "", kMusicStartToken); + add_special_token(vocab, "", kMusicEndToken); +} + +std::shared_ptr load_tiktoken_vocabulary(const std::filesystem::path & vocab_path) { + std::ifstream input(vocab_path, std::ios::binary); + if (!input) { + throw std::runtime_error("failed to open Yue2 tiktoken vocabulary: " + vocab_path.string()); + } + auto vocab = std::make_shared(); + vocab->pre_type = vendor::PreTokenizerType::Qwen2; + + std::string line; + int64_t mergeable_count = 0; + while (std::getline(input, line)) { + if (!line.empty() && line.back() == '\r') { + line.pop_back(); + } + if (line.empty()) { + continue; + } + std::istringstream parts(line); + std::string token_base64; + int64_t rank = -1; + if (!(parts >> token_base64 >> rank) || rank < 0 || rank > INT_MAX) { + throw std::runtime_error("Yue2 tiktoken vocabulary has an invalid line: " + line); + } + const std::string bytes = decode_base64(token_base64); + const auto token_id = static_cast(rank); + const std::string mapped = map_token_bytes(bytes); + vocab->token_to_id.emplace(mapped, token_id); + vocab->id_to_token.emplace(token_id, vendor::TokenData{mapped, 0}); + for (size_t split = 1; split < bytes.size(); ++split) { + vocab->bpe_ranks.emplace( + pair_key(map_token_bytes(bytes.substr(0, split)), map_token_bytes(bytes.substr(split))), + token_id); + } + ++mergeable_count; + } + if (mergeable_count != kEodToken) { + throw std::runtime_error("Yue2 tiktoken mergeable rank count must be 151643"); + } + register_yue2_special_tokens(*vocab, static_cast(mergeable_count)); + vendor::rebuild_special_tokens_cache(*vocab); + return vocab; +} + +} // namespace + +Yue2TextTokenizer::Yue2TextTokenizer(const std::filesystem::path & vocab_path) + : vocab_(load_tiktoken_vocabulary(vocab_path)) {} + +std::vector Yue2TextTokenizer::encode(const std::string & text) const { + return vendor::tokenize_bpe(*vocab_, text, true); +} + +} // namespace engine::models::yue2 diff --git a/src/models/yue2/types.cpp b/src/models/yue2/types.cpp new file mode 100644 index 00000000..3e662259 --- /dev/null +++ b/src/models/yue2/types.cpp @@ -0,0 +1,51 @@ +#include "engine/models/yue2/types.h" + +#include + +namespace engine::models::yue2 { + +const char * cot_mode_name(Yue2CotMode mode) noexcept { + switch (mode) { + case Yue2CotMode::Off: + return "off"; + case Yue2CotMode::Melody: + return "melody"; + case Yue2CotMode::Full: + return "full"; + } + return "full"; +} + +Yue2CotMode parse_cot_mode(const std::string & value) { + if (value == "off") { + return Yue2CotMode::Off; + } + if (value == "melody") { + return Yue2CotMode::Melody; + } + if (value == "full") { + return Yue2CotMode::Full; + } + throw std::runtime_error("yue2.cot must be one of off, melody, or full"); +} + +const char * cot_instruction(Yue2CotMode mode) noexcept { + switch (mode) { + case Yue2CotMode::Off: + return "Generate music with codec tokens from the given conditions."; + case Yue2CotMode::Melody: + return "Generate a melody-only ABC transcription without chord symbols, then generate music with codec tokens from the given conditions."; + case Yue2CotMode::Full: + return "Generate a chord-annotated ABC transcription, then generate music with codec tokens from the given conditions."; + } + return "Generate a chord-annotated ABC transcription, then generate music with codec tokens from the given conditions."; +} + +float request_guidance_scale(const Yue2Request & request) noexcept { + if (request.cfg_scale >= 0.0F) { + return request.cfg_scale; + } + return request.cot == Yue2CotMode::Off ? 1.01F : 1.0F; +} + +} // namespace engine::models::yue2 diff --git a/tests/unittests/test_qwen_decoder_packed_projections.cpp b/tests/unittests/test_qwen_decoder_packed_projections.cpp index 981063d4..25cd5735 100644 --- a/tests/unittests/test_qwen_decoder_packed_projections.cpp +++ b/tests/unittests/test_qwen_decoder_packed_projections.cpp @@ -1,4 +1,5 @@ #include "engine/framework/core/backend.h" +#include "engine/framework/debug/profiler.h" #include "engine/framework/modules/transformers/qwen_causal_decoder.h" #include "engine/framework/modules/transformers/qwen_decoder.h" #include "engine/framework/modules/optimizations/fast_kv_modules.h" @@ -7,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -36,6 +38,9 @@ void require_allclose( throw std::runtime_error(label + " size mismatch"); } for (size_t i = 0; i < actual.size(); ++i) { + if (!std::isfinite(actual[i]) || !std::isfinite(expected[i])) { + throw std::runtime_error(label + " contains non-finite values"); + } const float diff = std::fabs(actual[i] - expected[i]); if (diff > tolerance) { std::ostringstream message; @@ -52,9 +57,12 @@ struct LayerResult { std::vector value; }; -LayerResult run_layer(bool packed) { - constexpr int64_t batch = 1; - constexpr int64_t steps = 3; +LayerResult run_layer(bool packed, + engine::core::BackendType backend_type = engine::core::BackendType::Cpu, + bool batched_decode = false) { + const int64_t batch = batched_decode ? 2 : 1; + const int64_t steps = batched_decode ? 1 : 3; + constexpr int64_t cache_steps = 8; constexpr int64_t hidden = 8; constexpr int64_t heads = 2; constexpr int64_t kv_heads = 1; @@ -63,10 +71,10 @@ LayerResult run_layer(bool packed) { constexpr int64_t q_out = heads * head_dim; constexpr int64_t kv_out = kv_heads * head_dim; - engine::core::BackendConfig backend_config{engine::core::BackendType::Cpu, 0, 4}; + engine::core::BackendConfig backend_config{backend_type, 0, 8}; ggml_backend_t backend = engine::core::init_backend(backend_config); if (backend == nullptr) { - throw std::runtime_error("failed to initialize CPU backend"); + throw std::runtime_error("failed to initialize test backend"); } ggml_init_params params{kGraphBytes, nullptr, true}; @@ -78,7 +86,7 @@ LayerResult run_layer(bool packed) { ggml_backend_buffer_t buffer = nullptr; try { - engine::core::ModuleBuildContext ctx{ggml, "qwen_packed_projection_test", engine::core::BackendType::Cpu}; + engine::core::ModuleBuildContext ctx{ggml, "qwen_packed_projection_test", backend_type}; auto make_f32 = [&](std::initializer_list dims) { return engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims(dims)); }; @@ -87,7 +95,7 @@ LayerResult run_layer(bool packed) { auto positions = engine::core::make_tensor( ctx, GGML_TYPE_I32, - engine::core::TensorShape::from_dims({steps})); + engine::core::TensorShape::from_dims({batched_decode ? batch : steps})); engine::modules::QwenDecoderLayerWeights weights; weights.input_norm = {make_f32({hidden}), std::nullopt}; @@ -131,21 +139,63 @@ LayerResult run_layer(bool packed) { config.use_qk_norm = false; config.runtime.attention.prefill_mode = engine::modules::QwenDecoderAttentionMode::ManualRepeat; - const auto outputs = engine::modules::QwenDecoderLayerModule(config).build( - ctx, - input, - positions, - weights); - ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + engine::modules::QwenDecoderLayerOutputs outputs; + engine::core::TensorValue cache_key, cache_value, cache_slot, mask; + if (batched_decode) { + cache_key = make_f32({batch, cache_steps, kv_heads, head_dim}); + cache_value = make_f32({batch, cache_steps, kv_heads, head_dim}); + cache_slot = engine::core::make_tensor(ctx, GGML_TYPE_I32, + engine::core::TensorShape::from_dims({batch})); + mask = make_f32({batch, 1, 1, cache_steps}); + config.runtime.static_cache.update_mode = engine::modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + config.runtime.attention.static_mode = engine::modules::QwenDecoderAttentionMode::ManualRepeat; + outputs = engine::modules::QwenDecoderLayerModule(config).build_with_static_cache_tail_batched( + ctx, graph, input, positions, weights, cache_key, cache_value, cache_slot, mask); + } else { + outputs = engine::modules::QwenDecoderLayerModule(config).build(ctx, input, positions, weights); + } ggml_build_forward_expand(graph, outputs.output.tensor); + if (batched_decode) { + int rope_nodes = 0; + int copied_positions = 0; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + const auto * node = ggml_graph_node(graph, i); + if (node->op != GGML_OP_ROPE) { + continue; + } + ++rope_nodes; + if (node->src[1]->op == GGML_OP_CONT) { + ++copied_positions; + } + if (backend_type == engine::core::BackendType::Vulkan && node->src[1]->view_offs != 0) { + throw std::runtime_error("Vulkan batched RoPE retains an offset position view"); + } + } + const int expected_copies = backend_type == engine::core::BackendType::Vulkan ? 2 : 0; + if (rope_nodes != 4 || copied_positions != expected_copies) { + throw std::runtime_error("Unexpected batched RoPE position layout"); + } + } buffer = ggml_backend_alloc_ctx_tensors(ggml, backend); if (buffer == nullptr) { throw std::runtime_error("failed to allocate test tensors"); } engine::core::write_tensor_f32(input, patterned(static_cast(batch * steps * hidden), 2.1f, 0.20f)); - engine::core::write_tensor_i32(positions, {0, 1, 2}); + engine::core::write_tensor_i32(positions, batched_decode ? std::vector{2, 5} : std::vector{0, 1, 2}); + if (batched_decode) { + engine::core::write_tensor_f32(cache_key, patterned(batch * cache_steps * kv_out, 0.4f, 0.1f)); + engine::core::write_tensor_f32(cache_value, patterned(batch * cache_steps * kv_out, 0.8f, 0.1f)); + engine::core::write_tensor_i32(cache_slot, {2, cache_steps + 5}); + std::vector mask_values(batch * cache_steps, -std::numeric_limits::infinity()); + for (int64_t b = 0; b < batch; ++b) { + for (int64_t t = 0; t <= (b == 0 ? 2 : 5); ++t) { + mask_values[b * cache_steps + t] = 0.0f; + } + } + engine::core::write_tensor_f32(mask, mask_values); + } engine::core::write_tensor_f32(*weights.input_norm.weight, patterned(hidden, 0.3f, 0.7f)); engine::core::write_tensor_f32(*weights.post_norm.weight, patterned(hidden, 0.7f, 0.8f)); engine::core::write_tensor_f32( @@ -176,7 +226,9 @@ LayerResult run_layer(bool packed) { engine::core::write_tensor_f32(weights.mlp.up_proj.weight, up_values); } - ggml_backend_graph_compute(backend, graph); + if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("test graph execution failed"); + } LayerResult result; engine::core::read_tensor_f32_into(outputs.output.tensor, result.output); engine::core::read_tensor_f32_into(outputs.key.tensor, result.key); @@ -567,9 +619,29 @@ void test_higgs_decode_graph_exposes_cuda_fast_paths() { } // namespace -int main() { +int main(int argc, char ** argv) { try { + bool vulkan = false; + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "--vulkan") { + vulkan = true; + } else if (arg == "--log") { + engine::debug::configure_logging({true, ""}); + } else { + throw std::runtime_error("Unknown argument: " + arg); + } + } test_packed_qkv_and_gate_up_match_separate_projections(); + for (bool packed : {false, true}) { + const auto reference = run_layer(packed, engine::core::BackendType::Cpu, true); + if (vulkan) { + const auto actual = run_layer(packed, engine::core::BackendType::Vulkan, true); + require_allclose(actual.output, reference.output, 2.0e-5f, "Vulkan batched decoder output"); + require_allclose(actual.key, reference.key, 2.0e-5f, "Vulkan batched decoder key"); + require_allclose(actual.value, reference.value, 2.0e-5f, "Vulkan batched decoder value"); + } + } test_suffix_causal_mask(); test_f16_kv_set_rows(); test_f16_kv_set_rows_batched(); diff --git a/tests/unittests/test_sheetsage_audio_frontend.cpp b/tests/unittests/test_sheetsage_audio_frontend.cpp new file mode 100644 index 00000000..247f9498 --- /dev/null +++ b/tests/unittests/test_sheetsage_audio_frontend.cpp @@ -0,0 +1,60 @@ +#include "engine/models/sheetsage/audio_frontend.h" +#include "engine/framework/audio/wav_reader.h" + +#include +#include +#include +#include +#include +#include + +int main(int argc, char ** argv) { + try { + engine::models::sheetsage::SheetSage2AudioFrontend frontend; + const std::vector stereo{1, 0, 0, 1, 1, -1}; + const auto mono = frontend.prepare(stereo, 24000, 2, 24000, 8); + if (mono.size() != 3 || std::abs(mono[0] - std::sqrt(0.5F)) > 1e-7F || + mono[0] != mono[1] || mono[2] != 0) { + throw std::runtime_error("stereo mixing mismatch"); + } + for (int rate : {8000, 16000, 22050, 32000, 44100, 48000, 96000}) { + const std::vector constant(2001, 0.25F); + const auto first = frontend.prepare(constant, rate, 1, 24000, 8); + if (first != frontend.prepare(constant, rate, 1, 24000, 8)) { + throw std::runtime_error("cached filter mismatch"); + } + for (float value : first) { + if (std::abs(value - 0.25F) > 2e-5F) { + throw std::runtime_error("DC gain mismatch"); + } + } + } + // Optional external oracle pairs: input WAV followed by reference mono WAV. + for (int arg = 1; arg < argc; ++arg) { + if (std::string(argv[arg]) == "--log") continue; + if (arg + 1 >= argc) throw std::runtime_error("expected WAV pair"); + const auto input = engine::audio::read_wav_f32(std::filesystem::path(argv[arg])); + const auto reference = engine::audio::read_wav_f32(std::filesystem::path(argv[++arg])); + const auto actual = frontend.prepare(input.samples, input.sample_rate, input.channels, + reference.sample_rate, 8); + if (reference.channels != 1 || actual.size() != reference.samples.size()) { + throw std::runtime_error("oracle sample count mismatch: " + std::to_string(actual.size()) + + " vs " + std::to_string(reference.samples.size())); + } + double maximum = 0, squared = 0; + for (size_t i = 0; i < actual.size(); ++i) { + const double error = actual[i] - reference.samples[i]; + maximum = std::max(maximum, std::abs(error)); + squared += error * error; + } + std::cout << argv[arg - 1] << " samples=" << actual.size() << " max=" << maximum + << " rms=" << std::sqrt(squared / std::max(1, actual.size())) << '\n'; + if (maximum > 2e-6) throw std::runtime_error("oracle sample mismatch"); + } + std::cout << "SheetSage2 frontend passed\n"; + return 0; + } catch (const std::exception & error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/tests/unittests/test_sheetsage_processing.cpp b/tests/unittests/test_sheetsage_processing.cpp new file mode 100644 index 00000000..613d8410 --- /dev/null +++ b/tests/unittests/test_sheetsage_processing.cpp @@ -0,0 +1,47 @@ +#include "engine/models/sheetsage/processing.h" + +#include +#include +#include +#include +#include + +int main() { + try { + // Expected spellings from the Python notation implementation. + const std::vector> cases = { + {"A#:major", "A#:maj/3", "Bb", "A#/C##"}, + {"D#:major", "D#:maj7/7", "Eb", "D#maj7/C##"}, + {"G#:major", "G:min7/5", "Ab", "Gm7/D"}, + {"G:minor", "G:min/b3", "Gm", "Gm/Bb"}, + {"A#:minor", "A#:maj/5", "A#m", "A#/E#"}, + {"C:major", "C:maj/2", "C", "C/D"}, + {"G:major", "G:maj", "G", "G"}, + }; + for (const auto & test : cases) { + std::vector events(9); + for (size_t i = 0; i < events.size(); ++i) { + auto & event = events[i]; + event.time = static_cast(i) * 0.5F; + event.subbeat = event.source_subbeat = event.global_subbeat = static_cast(i) * 8; + event.timestamp = event.time; + event.meter = std::make_pair(4, 4); + event.eighth_position = static_cast(i % 4) * 2; + } + events[0].key = test[0]; + events[0].chord = test[1]; + events[0].notes.push_back({70, 0, 3, 4}); + events[0].note_end_times.push_back(0.5F); + const auto abc = engine::models::sheetsage::events_to_abc(events, 4.0); + if (abc.find("K:" + test[2] + "\n") == std::string::npos || + abc.find("\"" + test[3] + "\"") == std::string::npos) { + throw std::runtime_error("notation mismatch for " + test[0] + " " + test[1] + "\n" + abc); + } + } + std::cout << "PASS SheetSage2 key signatures and chord inversions\n"; + return 0; + } catch (const std::exception & error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/tests/yue2/sheetsage2_decoder_parity_probe.cpp b/tests/yue2/sheetsage2_decoder_parity_probe.cpp new file mode 100644 index 00000000..a4378638 --- /dev/null +++ b/tests/yue2/sheetsage2_decoder_parity_probe.cpp @@ -0,0 +1,144 @@ +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" +#include "engine/models/sheetsage/runtime.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +std::string require_arg(int argc, char ** argv, const std::string & name) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) { + return argv[i + 1]; + } + } + throw std::runtime_error("missing required argument " + name); +} + +std::string optional_arg(int argc, char ** argv, const std::string & name, const std::string & fallback) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) { + return argv[i + 1]; + } + } + return fallback; +} + +int64_t require_i64(int argc, char ** argv, const std::string & name) { + return std::stoll(require_arg(argc, argv, name)); +} + +engine::core::BackendType parse_backend(const std::string & value) { + if (value == "cpu") return engine::core::BackendType::Cpu; + if (value == "cuda") return engine::core::BackendType::Cuda; + if (value == "vulkan") return engine::core::BackendType::Vulkan; + if (value == "metal") return engine::core::BackendType::Metal; + throw std::runtime_error("unsupported backend: " + value); +} + +template +std::vector read_binary(const std::filesystem::path & path) { + std::ifstream in(path, std::ios::binary); + if (!in) { + throw std::runtime_error("failed to open input file: " + path.string()); + } + in.seekg(0, std::ios::end); + const std::streamoff bytes = in.tellg(); + in.seekg(0, std::ios::beg); + if (bytes < 0 || bytes % static_cast(sizeof(T)) != 0) { + throw std::runtime_error("input byte size is not aligned: " + path.string()); + } + std::vector values(static_cast(bytes / static_cast(sizeof(T)))); + in.read(reinterpret_cast(values.data()), bytes); + if (!in) { + throw std::runtime_error("failed to read input file: " + path.string()); + } + return values; +} + +void write_f32(const std::filesystem::path & path, const std::vector & values) { + std::ofstream out(path, std::ios::binary); + if (!out) { + throw std::runtime_error("failed to open output file: " + path.string()); + } + out.write(reinterpret_cast(values.data()), static_cast(values.size() * sizeof(float))); +} + +std::filesystem::path resolve_tensor_file(const std::filesystem::path & model) { + if (std::filesystem::is_directory(model)) { + const auto direct = model / "model.safetensors"; + if (std::filesystem::is_regular_file(direct)) { + return direct; + } + } + return model; +} + +void compare(const std::vector & got, const std::vector & ref) { + if (got.size() != ref.size()) { + throw std::runtime_error("output/reference size mismatch"); + } + double dot = 0.0; + double got_norm = 0.0; + double ref_norm = 0.0; + double mse = 0.0; + float max_abs = 0.0F; + for (size_t i = 0; i < got.size(); ++i) { + const double g = got[i]; + const double r = ref[i]; + const double d = g - r; + dot += g * r; + got_norm += g * g; + ref_norm += r * r; + mse += d * d; + max_abs = std::max(max_abs, static_cast(std::abs(d))); + } + const double rmse = std::sqrt(mse / std::max(got.size(), 1)); + const double cosine = dot / std::sqrt(std::max(got_norm * ref_norm, 1.0e-30)); + std::cout << "output_values=" << got.size() << "\n"; + std::cout << "max_abs=" << max_abs << "\n"; + std::cout << "rmse=" << rmse << "\n"; + std::cout << "cosine=" << cosine << "\n"; +} + +} // namespace + +int main(int argc, char ** argv) { + try { + const auto model = std::filesystem::path(require_arg(argc, argv, "--model")); + const auto mixed_path = std::filesystem::path(require_arg(argc, argv, "--mixed")); + const auto ids_path = std::filesystem::path(require_arg(argc, argv, "--ids")); + const auto output_path = std::filesystem::path(optional_arg(argc, argv, "--output", "")); + const auto reference_path = std::filesystem::path(optional_arg(argc, argv, "--reference", "")); + const int64_t memory_steps = require_i64(argc, argv, "--memory-steps"); + const int threads = static_cast(std::stoll(optional_arg(argc, argv, "--threads", "8"))); + engine::core::ExecutionContext execution({parse_backend(optional_arg(argc, argv, "--backend", "cpu")), 0, threads}); + auto source = engine::assets::open_tensor_source(resolve_tensor_file(model)); + engine::models::sheetsage::SheetSage2DecoderRuntime runtime(source, execution); + const auto mixed = read_binary(mixed_path); + const auto ids = read_binary(ids_path); + const auto logits = runtime.decode_logits(mixed, memory_steps, ids); + if (!output_path.empty()) { + write_f32(output_path, logits); + } + if (!reference_path.empty()) { + compare(logits, read_binary(reference_path)); + } else { + std::cout << "output_values=" << logits.size() << "\n"; + } + return 0; + } catch (const std::exception & e) { + std::cerr << "sheetsage2_decoder_parity_probe failed: " << e.what() << "\n"; + return 1; + } +} diff --git a/tests/yue2/sheetsage2_reference_dump.py b/tests/yue2/sheetsage2_reference_dump.py new file mode 100644 index 00000000..24be20e0 --- /dev/null +++ b/tests/yue2/sheetsage2_reference_dump.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 + +import argparse +import importlib +import json +import sys +from pathlib import Path + +import numpy as np +import torch + + +def write_f32(path: Path, tensor: torch.Tensor) -> None: + path.write_bytes(tensor.detach().cpu().contiguous().float().numpy().astype(" None: + path.write_bytes(tensor.detach().cpu().contiguous().to(torch.int32).numpy().astype(" None: + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True) + parser.add_argument("--base-model", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--audio-frames", type=int, default=24000 * 8) + parser.add_argument("--decoder-steps", type=int, default=48) + parser.add_argument("--seed", type=int, default=20260910) + args = parser.parse_args() + + root = Path(__file__).resolve().parents[2] + sys.path.insert(0, str((root / "reference").resolve())) + transformers = importlib.import_module("transformers452") + sys.modules["transformers"] = transformers + for name in ( + "configuration_utils", + "dynamic_module_utils", + "generation", + "generation.utils", + "modeling_outputs", + "modeling_utils", + "models", + "models.auto", + "models.auto.configuration_auto", + "models.auto.modeling_auto", + "models.bart", + "models.bart.configuration_bart", + "models.bart.modeling_bart", + "utils", + "utils.hub", + ): + try: + sys.modules[f"transformers.{name}"] = importlib.import_module(f"transformers452.{name}") + except ModuleNotFoundError: + pass + import transformers452.dynamic_module_utils as transformers452_dynamic_module_utils + transformers452_dynamic_module_utils.transformers = transformers + from SheetSage2.modeling_sheetsage2 import SheetSage2Model + + torch.manual_seed(args.seed) + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + model = SheetSage2Model.from_pretrained( + args.model, + base_model_path=args.base_model, + local_files_only=True, + torch_dtype=torch.float32, + device_map="cpu", + ).eval() + + # A deterministic full-band synthetic waveform is long enough to exercise + # the real MERT2 front end without relying on a tiny smoke shape. + t = torch.arange(args.audio_frames, dtype=torch.float32) / float(model.config.sampling_rate) + waveform = ( + 0.20 * torch.sin(2.0 * torch.pi * 220.0 * t) + + 0.07 * torch.sin(2.0 * torch.pi * 440.0 * t + 0.3) + + 0.03 * torch.sin(2.0 * torch.pi * 880.0 * t + 0.9) + ).unsqueeze(0) + + with torch.no_grad(): + enc = model.get_audio_features(waveform, output_hidden_states=False, return_dict=True) + mixed = enc.mixed_hidden_state.detach().contiguous() + ids = torch.arange(args.decoder_steps, dtype=torch.long).unsqueeze(0) + ids = (ids * 17 + 5) % model.config.vocab_size + ids[:, 0] = model.config.bos_token_id + logits, _ = model.decode(mixed @ model.encoder_projection.weight.T + model.encoder_projection.bias, ids) + + write_f32(out_dir / "mixed_encoder_state.f32", mixed) + write_i32(out_dir / "decoder_input_ids.i32", ids) + write_f32(out_dir / "logits_ref.f32", logits) + meta = { + "seed": args.seed, + "audio_frames": args.audio_frames, + "decoder_steps": args.decoder_steps, + "batch": int(mixed.shape[0]), + "memory_steps": int(mixed.shape[1]), + "encoder_hidden_size": int(mixed.shape[2]), + "logits_shape": [int(v) for v in logits.shape], + } + (out_dir / "meta.json").write_text(json.dumps(meta, indent=2) + "\n", encoding="utf-8") + print(json.dumps(meta, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tests/yue2/yue2_python_warm_bench.py b/tests/yue2/yue2_python_warm_bench.py new file mode 100644 index 00000000..45ae7bc9 --- /dev/null +++ b/tests/yue2/yue2_python_warm_bench.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import json +import os +import random +import sys +import time +from pathlib import Path +from typing import Any + +import numpy as np +import soundfile as sf +import torch + + +REPO_ROOT = Path(__file__).resolve().parents[2] +REFERENCE_ROOT = REPO_ROOT / "reference" / "YuE" +DEFAULT_CASES = REPO_ROOT / "tests" / "yue2" / "yue2_warm_bench_cases.json" +DEFAULT_OUTPUT_ROOT = REPO_ROOT / "build" / "logs" / "yue2" / "python_baseline" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Python reference YuE2 warmbench.") + parser.add_argument("--family", default="yue2") + parser.add_argument("--model", default="/home/leo/Desktop/YuE2/YuE2-3B") + parser.add_argument("--vae", default="/home/leo/Desktop/YuE2/YuE2-Vae") + parser.add_argument("--reference-root", type=Path, default=REFERENCE_ROOT) + parser.add_argument("--cases", type=Path, default=DEFAULT_CASES) + parser.add_argument("--case", action="append", default=[]) + parser.add_argument("--backend", choices=("cuda", "cpu"), default="cuda") + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--threads", type=int, default=8) + parser.add_argument("--warmup", type=int, default=0) + parser.add_argument("--iterations", type=int, default=1) + parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_ROOT) + parser.add_argument("--timing-file", type=Path, default=DEFAULT_OUTPUT_ROOT / "python_baseline.log") + parser.add_argument("--summary-file", type=Path, default=DEFAULT_OUTPUT_ROOT / "summary.json") + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument("--log", action="store_true") + return parser.parse_args() + + +def resolve_path(path: Path | str) -> Path: + value = Path(path) + return value if value.is_absolute() else REPO_ROOT / value + + +def add_reference_path(reference_root: Path) -> None: + src = resolve_path(reference_root) / "src" + if not (src / "yue2" / "__init__.py").is_file(): + raise RuntimeError(f"missing YuE2 reference package under {src}") + sys.path.insert(0, str(src.resolve())) + + +def configure_runtime(args: argparse.Namespace) -> str: + os.environ.setdefault("PYTHONHASHSEED", "0") + torch.set_num_threads(max(1, args.threads)) + random.seed(0) + np.random.seed(0) + torch.manual_seed(0) + if args.backend == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("YuE2 warmbench requested CUDA, but torch.cuda.is_available() is false") + torch.cuda.set_device(args.device) + torch.cuda.manual_seed_all(0) + return f"cuda:{args.device}" + return "cpu" + + +def sync_device(device: str) -> None: + if device.startswith("cuda") and torch.cuda.is_available(): + torch.cuda.synchronize() + + +def load_cases(path: Path) -> dict[str, dict[str, Any]]: + payload = json.loads(resolve_path(path).read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise RuntimeError("YuE2 warmbench cases must be a JSON object") + return payload + + +def selected_cases(args: argparse.Namespace, cases: dict[str, dict[str, Any]]) -> list[dict[str, Any]]: + names = args.case or list(cases) + out = [] + for name in names: + if name not in cases: + raise RuntimeError(f"unknown YuE2 warmbench case: {name}") + case = dict(cases[name]) + case.setdefault("id", name) + out.append(case) + return out + + +def make_generation_config(case: dict[str, Any]): + from yue2.protocol import GenerationConfig, Sampling + + defaults = GenerationConfig() + abc = Sampling( + temperature=float(case.get("abc_temperature", defaults.abc.temperature)), + top_p=float(case.get("abc_top_p", defaults.abc.top_p)), + top_k=int(case.get("abc_top_k", defaults.abc.top_k)), + repetition_penalty=float(case.get("abc_repetition_penalty", defaults.abc.repetition_penalty)), + penalty_window=int(case.get("abc_penalty_window", defaults.abc.penalty_window)), + min_tokens=int(case.get("abc_min_tokens", defaults.abc.min_tokens)), + max_tokens=int(case.get("abc_max_tokens", defaults.abc.max_tokens)), + ) + semantic = Sampling( + temperature=float(case.get("semantic_temperature", defaults.semantic.temperature)), + top_p=float(case.get("semantic_top_p", defaults.semantic.top_p)), + top_k=int(case.get("semantic_top_k", defaults.semantic.top_k)), + repetition_penalty=float(case.get("semantic_repetition_penalty", defaults.semantic.repetition_penalty)), + penalty_window=int(case.get("semantic_penalty_window", defaults.semantic.penalty_window)), + min_tokens=int(case.get("semantic_min_tokens", defaults.semantic.min_tokens)), + max_tokens=int(case.get("semantic_max_tokens", defaults.semantic.max_tokens)), + ) + return GenerationConfig( + abc=abc, + semantic=semantic, + ode_steps=int(case.get("ode_steps", defaults.ode_steps)), + ode_method=str(case.get("ode_method", defaults.ode_method)), + context=int(case.get("context", defaults.context)), + version=str(case.get("version", defaults.version)), + ) + + +def request_from_case(case: dict[str, Any]) -> dict[str, Any]: + request = { + "style": str(case["style"]), + "lyrics": str(case["lyrics"]), + "cot": str(case.get("cot", "full")), + "seed": int(case.get("seed", 831001)), + "id": str(case.get("id", "song")), + } + if "cfg_scale" in case: + request["cfg_scale"] = float(case["cfg_scale"]) + if "abc" in case: + request["abc"] = str(case["abc"]) + if "abc_file" in case: + request["abc"] = resolve_path(case["abc_file"]).read_text(encoding="utf-8") + return request + + +def summarize_audio(audio: np.ndarray, sample_rate: int) -> dict[str, Any]: + flat = np.asarray(audio, dtype=np.float32).reshape(-1) + if flat.size == 0: + raise RuntimeError("YuE2 warmbench produced empty audio") + return { + "sample_rate": int(sample_rate), + "channels": int(audio.shape[1]) if audio.ndim == 2 else 1, + "frames": int(audio.shape[0]) if audio.ndim >= 1 else 0, + "samples": int(flat.size), + "duration_sec": float((audio.shape[0] if audio.ndim >= 1 else 0) / sample_rate), + "sum": float(np.sum(flat, dtype=np.float64)), + "mean_abs": float(np.mean(np.abs(flat), dtype=np.float64)), + "rms": float(np.sqrt(np.mean(np.square(flat), dtype=np.float64))), + "min": float(np.min(flat)), + "max": float(np.max(flat)), + } + + +def run_case(pipe: Any, case: dict[str, Any], iteration: int, output_root: Path) -> dict[str, Any]: + request = request_from_case(case) + case_dir = output_root / f"{request['id']}_iter{iteration}" + case_dir.mkdir(parents=True, exist_ok=True) + config = make_generation_config(case) + pipe.generation_config = config + + start = time.perf_counter() + plan_start = time.perf_counter() + plan = pipe.plan(**request) + sync_device(str(pipe.device)) + plan_wall = time.perf_counter() - plan_start + + semantic_start = time.perf_counter() + semantic = pipe.generate_semantic(plan) + sync_device(str(pipe.device)) + semantic_wall = time.perf_counter() - semantic_start + + nar_start = time.perf_counter() + latents = pipe.synthesize(semantic) + sync_device(str(pipe.device)) + nar_wall = time.perf_counter() - nar_start + + decode_start = time.perf_counter() + audio = pipe.decode(latents) + sync_device(str(pipe.device)) + decode_wall = time.perf_counter() - decode_start + wall = time.perf_counter() - start + + audio_path = case_dir / "audio.wav" + sf.write(audio_path, audio, 48000, subtype="FLOAT") + plan.save(case_dir / "plan") + np.save(case_dir / "semantic.npy", np.asarray(semantic.tokens, dtype=np.int32)) + np.save(case_dir / "latent.npy", latents.astype(np.float32)) + + result = { + "case": request["id"], + "iteration": iteration, + "request": request, + "truncated": {"abc": bool(plan.truncated), "semantic": bool(semantic.truncated)}, + "token_counts": { + "abc": len(plan.abc_ids), + "prefix": len(plan.prefix), + "semantic": len(semantic.tokens), + "latent_frames": int(latents.shape[0]) if getattr(latents, "ndim", 0) == 2 else None, + }, + "timing_sec": { + "plan_wall": plan_wall, + "semantic_wall": semantic_wall, + "nar_wall": nar_wall, + "decode_wall": decode_wall, + "wall": wall, + }, + "audio": summarize_audio(audio, 48000), + "paths": { + "case_dir": str(case_dir), + "audio": str(audio_path), + }, + } + (case_dir / "result.json").write_text(json.dumps(result, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") + return result + + +def main() -> int: + args = parse_args() + add_reference_path(args.reference_root) + device = configure_runtime(args) + from yue2 import YuE2Pipeline + + output_root = resolve_path(args.output_dir) + output_root.mkdir(parents=True, exist_ok=True) + args.timing_file = resolve_path(args.timing_file) + args.summary_file = resolve_path(args.summary_file) + args.timing_file.parent.mkdir(parents=True, exist_ok=True) + args.summary_file.parent.mkdir(parents=True, exist_ok=True) + + cases = selected_cases(args, load_cases(args.cases)) + results: list[dict[str, Any]] = [] + model = str(resolve_path(args.model)) if Path(args.model).exists() else args.model + vae = str(resolve_path(args.vae)) if Path(args.vae).exists() else args.vae + + with YuE2Pipeline.from_pretrained( + model, + vae=vae, + device=device, + backend="torch", + local_files_only=args.local_files_only, + progress=False, + generation_config=make_generation_config(cases[0]), + ) as pipe: + for i in range(args.warmup): + run_case(pipe, cases[0], -(i + 1), output_root / "warmup") + for iteration in range(args.iterations): + for case in cases: + result = run_case(pipe, case, iteration, output_root) + results.append(result) + with args.timing_file.open("a", encoding="utf-8") as log: + log.write(json.dumps(result, ensure_ascii=False, sort_keys=True) + "\n") + + summary = { + "family": args.family, + "model": model, + "vae": vae, + "backend": args.backend, + "device": device, + "threads": args.threads, + "cases": [r["case"] for r in results], + "results": results, + } + args.summary_file.write_text(json.dumps(summary, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/yue2/yue2_vae_parity_probe.cpp b/tests/yue2/yue2_vae_parity_probe.cpp new file mode 100644 index 00000000..4a97b7df --- /dev/null +++ b/tests/yue2/yue2_vae_parity_probe.cpp @@ -0,0 +1,148 @@ +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/codecs/oobleck_audio_vae_runtime.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +std::string arg_value(int argc, char ** argv, const std::string & name, const std::string & fallback = {}) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) { + return argv[i + 1]; + } + } + return fallback; +} + +int64_t int_arg(int argc, char ** argv, const std::string & name, int64_t fallback) { + return std::stoll(arg_value(argc, argv, name, std::to_string(fallback))); +} + +engine::core::BackendType parse_backend(const std::string & value) { + if (value == "cpu") return engine::core::BackendType::Cpu; + if (value == "cuda") return engine::core::BackendType::Cuda; + if (value == "vulkan") return engine::core::BackendType::Vulkan; + throw std::runtime_error("unsupported backend: " + value); +} + +std::vector read_f32_file(const std::filesystem::path & path) { + std::ifstream in(path, std::ios::binary); + if (!in) { + throw std::runtime_error("failed to open " + path.string()); + } + in.seekg(0, std::ios::end); + const auto bytes = in.tellg(); + if (bytes < 0 || bytes % static_cast(sizeof(float)) != 0) { + throw std::runtime_error("invalid f32 byte size: " + path.string()); + } + in.seekg(0, std::ios::beg); + std::vector values(static_cast(bytes / static_cast(sizeof(float)))); + in.read(reinterpret_cast(values.data()), bytes); + return values; +} + +void write_f32_file(const std::filesystem::path & path, const std::vector & values) { + std::filesystem::create_directories(path.parent_path()); + std::ofstream out(path, std::ios::binary); + if (!out) { + throw std::runtime_error("failed to write " + path.string()); + } + out.write(reinterpret_cast(values.data()), static_cast(values.size() * sizeof(float))); +} + +struct CompareMetrics { + double max_abs = 0.0; + double rmse = 0.0; + double cosine = 0.0; +}; + +CompareMetrics compare(const std::vector & got, const std::vector & ref) { + if (got.size() != ref.size()) { + throw std::runtime_error("output/reference size mismatch"); + } + double sum_sq = 0.0; + double dot = 0.0; + double got_sq = 0.0; + double ref_sq = 0.0; + double max_abs = 0.0; + for (size_t i = 0; i < got.size(); ++i) { + const double g = got[i]; + const double r = ref[i]; + const double d = g - r; + max_abs = std::max(max_abs, std::abs(d)); + sum_sq += d * d; + dot += g * r; + got_sq += g * g; + ref_sq += r * r; + } + return { + max_abs, + std::sqrt(sum_sq / static_cast(got.size())), + dot / std::sqrt(std::max(got_sq * ref_sq, std::numeric_limits::min())), + }; +} + +} // namespace + +int main(int argc, char ** argv) { + try { + const auto mode = arg_value(argc, argv, "--mode", "decode"); + const auto model = std::filesystem::path(arg_value(argc, argv, "--model", "/home/leo/Desktop/YuE2/YuE2-Vae")); + const auto input_path = std::filesystem::path(arg_value(argc, argv, "--input")); + const auto output_path = std::filesystem::path(arg_value(argc, argv, "--output")); + const auto reference_path = std::filesystem::path(arg_value(argc, argv, "--reference")); + const int threads = static_cast(int_arg(argc, argv, "--threads", 8)); + const int64_t frames = int_arg(argc, argv, "--frames", 32); + const int64_t batch = int_arg(argc, argv, "--batch", 1); + + engine::core::BackendConfig backend_config; + backend_config.type = parse_backend(arg_value(argc, argv, "--backend", "cpu")); + backend_config.threads = threads; + engine::core::ExecutionContext execution(backend_config); + + const auto tensor_source_path = std::filesystem::is_directory(model) ? model / "model.safetensors" : model; + auto source = engine::assets::open_tensor_source(tensor_source_path); + engine::codecs::OobleckAudioVaeRuntime runtime(source, execution); + const auto input = read_f32_file(input_path); + std::vector output; + if (mode == "decode") { + output = runtime.decode_planar(input, batch, frames); + } else if (mode == "encode") { + output = runtime.encode_planar(input, frames); + } else { + throw std::runtime_error("unsupported mode: " + mode); + } + if (!output_path.empty()) { + write_f32_file(output_path, output); + } + std::cout << "mode=" << mode << " output_values=" << output.size() << "\n"; + if (!reference_path.empty()) { + const auto ref = read_f32_file(reference_path); + const auto m = compare(output, ref); + std::cout << "max_abs=" << m.max_abs << "\n"; + std::cout << "rmse=" << m.rmse << "\n"; + std::cout << "cosine=" << m.cosine << "\n"; + const double max_abs_limit = mode == "decode" ? 8.0e-4 : 1.5e-3; + const double cosine_limit = 0.999999; + if (m.max_abs > max_abs_limit || m.cosine < cosine_limit) { + throw std::runtime_error("YuE2 VAE parity check failed"); + } + } + return 0; + } catch (const std::exception & error) { + std::cerr << "yue2_vae_parity_probe failed: " << error.what() << "\n"; + return 1; + } +} diff --git a/tests/yue2/yue2_vae_reference_dump.py b/tests/yue2/yue2_vae_reference_dump.py new file mode 100644 index 00000000..e5ceb78a --- /dev/null +++ b/tests/yue2/yue2_vae_reference_dump.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import torch + + +def import_yue(reference_root: Path) -> None: + sys.path.insert(0, str(reference_root / "src")) + + +def write_f32(path: Path, tensor: torch.Tensor) -> None: + array = tensor.detach().cpu().contiguous().float().numpy() + path.parent.mkdir(parents=True, exist_ok=True) + array.tofile(path) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Dump realistic YuE2 VAE Python reference tensors.") + parser.add_argument("--reference-root", type=Path, default=Path("reference/YuE")) + parser.add_argument("--model", type=Path, default=Path("/home/leo/Desktop/YuE2/YuE2-Vae")) + parser.add_argument("--out-dir", type=Path, required=True) + parser.add_argument("--latent-frames", type=int, default=32) + parser.add_argument("--audio-frames", type=int, default=61440) + parser.add_argument("--seed", type=int, default=20260909) + args = parser.parse_args() + + import_yue(args.reference_root.resolve()) + from yue2.modeling_vae import YuE2VAE + + torch.manual_seed(args.seed) + torch.set_grad_enabled(False) + model = YuE2VAE.from_pretrained(args.model, device="cpu", decoder_only=False).eval() + + latents = torch.randn(1, model.config.latent_dim, args.latent_frames, dtype=torch.float32) * 0.35 + decoded = model.decoder(latents) + audio = torch.randn(1, model.config.audio_channels, args.audio_frames, dtype=torch.float32) * 0.08 + encoded = model.encoder(audio) + + args.out_dir.mkdir(parents=True, exist_ok=True) + write_f32(args.out_dir / "decode_input.f32", latents) + write_f32(args.out_dir / "decode_ref.f32", decoded) + write_f32(args.out_dir / "encode_input.f32", audio) + write_f32(args.out_dir / "encode_ref.f32", encoded) + metadata = { + "seed": args.seed, + "latent_frames": args.latent_frames, + "audio_frames": args.audio_frames, + "decode_shape": list(decoded.shape), + "encode_shape": list(encoded.shape), + "sample_rate": model.config.sample_rate, + "downsampling_ratio": model.config.downsampling_ratio, + } + (args.out_dir / "metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") + print(json.dumps(metadata, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tests/yue2/yue2_warm_bench_cases.json b/tests/yue2/yue2_warm_bench_cases.json new file mode 100644 index 00000000..477f2a39 --- /dev/null +++ b/tests/yue2/yue2_warm_bench_cases.json @@ -0,0 +1,56 @@ +{ + "direct_off_medium": { + "id": "direct_off_medium", + "style": "English, indie pop, bright acoustic guitar, soft drums, warm lead vocal, polished demo mix", + "lyrics": "[Verse]\nAudio dot cpp starts the demo tonight.\nSmall sparks of music glow in the light.\nA local engine keeps the rhythm tight.\nNo cloud in the loop, just code taking flight.\n[Pre-Chorus]\nEvery buffer finds its place.\nEvery model joins the race.\n[Chorus]\nTurn the signal into song.\nLet the native runtime carry it along.\nFrom text to melody, clear and strong.\nAudio dot cpp keeps the demo moving on.", + "cot": "off", + "seed": 20260910, + "cfg_scale": 1.01, + "abc_max_tokens": 1, + "abc_min_tokens": 0, + "semantic_max_tokens": 640, + "semantic_min_tokens": 192, + "ode_steps": 8 + }, + "create_full_plan_medium": { + "id": "create_full_plan_medium", + "style": "English, piano pop, clear lead vocal, gentle bass, 92 BPM, warm chorus harmonies, soft room reverb", + "lyrics": "[Verse]\nA quiet room begins to sing.\nThe keys reply with silver rings.\nA metronome is counting time.\nA simple phrase becomes a line.\n[Pre-Chorus]\nThe bass walks in with patient grace.\nThe melody opens more space.\n[Chorus]\nHold the note and let it rise.\nMorning opens up the sky.\nEvery word can find a place.\nEvery phrase can leave a trace.\n[Bridge]\nIf the plan is clear and bright.\nThe song can travel through the night.", + "cot": "full", + "seed": 20260911, + "cfg_scale": 1.0, + "abc_max_tokens": 512, + "abc_min_tokens": 128, + "semantic_max_tokens": 768, + "semantic_min_tokens": 224, + "ode_steps": 8 + }, + "cover_melody_score_medium": { + "id": "cover_melody_score_medium", + "style": "English, jazz funk cover, warm Rhodes, round bass, light drums, relaxed vocal, clean live band feel", + "lyrics": "[Verse]\nWe follow the melody line.\nThe rhythm keeps everything fine.\nA Rhodes chord lands behind the beat.\nThe bass keeps walking down the street.\n[Chorus]\nHold that shape and make it new.\nChange the color, keep the view.\nEvery phrase can still belong.\nWhen the cover turns into a song.", + "cot": "melody", + "abc_file": "reference/YuE/examples/melody.abc", + "seed": 20260912, + "cfg_scale": 1.0, + "abc_max_tokens": 1, + "abc_min_tokens": 0, + "semantic_max_tokens": 640, + "semantic_min_tokens": 192, + "ode_steps": 8 + }, + "edit_full_score_medium": { + "id": "edit_full_score_medium", + "style": "English, piano pop with jazz harmony, clear lead vocal, gentle bass, 92 BPM, wider chorus, tasteful drum fills", + "lyrics": "[Verse]\nA quiet room begins to sing.\nThe keys reply with silver rings.\nA brushed snare paints the second line.\nThe harmony turns right on time.\n[Chorus]\nHold the note and let it rise.\nMorning opens up the sky.\nKeep the score but change the light.\nMake the ending open wide.\n[Outro]\nLet the last chord slowly fade.\nLeave the echo we have made.", + "cot": "full", + "abc_file": "reference/YuE/examples/score-jazz.abc", + "seed": 20260913, + "cfg_scale": 1.0, + "abc_max_tokens": 1, + "abc_min_tokens": 0, + "semantic_max_tokens": 640, + "semantic_min_tokens": 192, + "ode_steps": 8 + } +} diff --git a/tools/audiocpp_cli/audiocpp_cli_path_cases.json b/tools/audiocpp_cli/audiocpp_cli_path_cases.json index deeef389..91ca9d0b 100644 --- a/tools/audiocpp_cli/audiocpp_cli_path_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_path_cases.json @@ -1106,6 +1106,120 @@ } ] }, + { + "id": "yue2_direct_off_medium", + "coverage": "YuE2 direct text-to-music route without ABC planning", + "family": "yue2", + "model": "models/YuE2-3B-Q8_0-GGUF", + "task": "gen", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "direct_off_medium", + "text": "[Verse]\nAudio dot cpp starts the demo tonight.\nSmall sparks of music glow in the light.\nA local engine keeps the rhythm tight.\nNo cloud in the loop, just code taking flight.\n[Pre-Chorus]\nEvery buffer finds its place.\nEvery model joins the race.\n[Chorus]\nTurn the signal into song.\nLet the native runtime carry it along.\nFrom text to melody, clear and strong.\nAudio dot cpp keeps the demo moving on.", + "seed": 20260910, + "num_inference_steps": 8, + "options": { + "cfg_scale": 1.01, + "style": "English, indie pop, bright acoustic guitar, soft drums, warm lead vocal, polished demo mix", + "cot": "off", + "abc_min_tokens": 0, + "abc_max_tokens": 1, + "semantic_min_tokens": 192, + "semantic_max_tokens": 640 + } + } + ] + }, + { + "id": "yue2_create_full_plan_medium", + "coverage": "YuE2 full ABC planning route", + "family": "yue2", + "model": "models/YuE2-3B-Q8_0-GGUF", + "task": "gen", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "create_full_plan_medium", + "text": "[Verse]\nA quiet room begins to sing.\nThe keys reply with silver rings.\nA metronome is counting time.\nA simple phrase becomes a line.\n[Pre-Chorus]\nThe bass walks in with patient grace.\nThe melody opens more space.\n[Chorus]\nHold the note and let it rise.\nMorning opens up the sky.\nEvery word can find a place.\nEvery phrase can leave a trace.\n[Bridge]\nIf the plan is clear and bright.\nThe song can travel through the night.", + "seed": 20260911, + "num_inference_steps": 8, + "options": { + "cfg_scale": 1.0, + "style": "English, piano pop, clear lead vocal, gentle bass, 92 BPM, warm chorus harmonies, soft room reverb", + "cot": "full", + "abc_min_tokens": 128, + "abc_max_tokens": 512, + "semantic_min_tokens": 224, + "semantic_max_tokens": 768 + } + } + ] + }, + { + "id": "yue2_cover_melody_score_medium", + "coverage": "YuE2 melody score conditioning route with external ABC", + "family": "yue2", + "model": "models/YuE2-3B-Q8_0-GGUF", + "task": "gen", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "cover_melody_score_medium", + "text": "[Verse]\nWe follow the melody line.\nThe rhythm keeps everything fine.\nA Rhodes chord lands behind the beat.\nThe bass keeps walking down the street.\n[Chorus]\nHold that shape and make it new.\nChange the color, keep the view.\nEvery phrase can still belong.\nWhen the cover turns into a song.", + "seed": 20260912, + "num_inference_steps": 8, + "options": { + "cfg_scale": 1.0, + "style": "English, jazz funk cover, warm Rhodes, round bass, light drums, relaxed vocal, clean live band feel", + "cot": "melody", + "abc_file": "reference/YuE/examples/melody.abc", + "abc_min_tokens": 0, + "abc_max_tokens": 1, + "semantic_min_tokens": 192, + "semantic_max_tokens": 640 + } + } + ] + }, + { + "id": "yue2_edit_full_score_medium", + "coverage": "YuE2 full score edit route with external ABC", + "family": "yue2", + "model": "models/YuE2-3B-Q8_0-GGUF", + "task": "gen", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "edit_full_score_medium", + "text": "[Verse]\nA quiet room begins to sing.\nThe keys reply with silver rings.\nA brushed snare paints the second line.\nThe harmony turns right on time.\n[Chorus]\nHold the note and let it rise.\nMorning opens up the sky.\nKeep the score but change the light.\nMake the ending open wide.\n[Outro]\nLet the last chord slowly fade.\nLeave the echo we have made.", + "seed": 20260913, + "num_inference_steps": 8, + "options": { + "cfg_scale": 1.0, + "style": "English, piano pop with jazz harmony, clear lead vocal, gentle bass, 92 BPM, wider chorus, tasteful drum fills", + "cot": "full", + "abc_file": "reference/YuE/examples/score-jazz.abc", + "abc_min_tokens": 0, + "abc_max_tokens": 1, + "semantic_min_tokens": 192, + "semantic_max_tokens": 640 + } + } + ] + }, { "id": "irodori_tts_500m_emoji_style_clone", "coverage": "Irodori-TTS 500M no-reference emoji/style requests and reference clone path in one session", diff --git a/tools/community_models/sheetsage2/convert_sheetsage2_gguf.py b/tools/community_models/sheetsage2/convert_sheetsage2_gguf.py new file mode 100644 index 00000000..4df64bf3 --- /dev/null +++ b/tools/community_models/sheetsage2/convert_sheetsage2_gguf.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +"""Convert SheetSage2 adapter + MERT2 parent into a self-contained GGUF.""" + +from __future__ import annotations + +import argparse +import json +import shutil +import subprocess +import tempfile +from pathlib import Path + +import torch +from safetensors.torch import save_file, safe_open + + +REPO_ROOT = Path(__file__).resolve().parents[3] +PROJECTIONS = ("query_proj", "key_proj", "value_proj", "out_proj") + + +def require_file(path: Path) -> Path: + if not path.is_file(): + raise FileNotFoundError(path) + return path + + +def load_safetensors(path: Path) -> dict[str, torch.Tensor]: + with safe_open(path, framework="pt", device="cpu") as handle: + return {key: handle.get_tensor(key) for key in handle.keys()} + + +def merge_adapter(adapter_dir: Path, mert_dir: Path, output_path: Path) -> None: + adapter_weights = load_safetensors(require_file(adapter_dir / "model.safetensors")) + mert_weights = load_safetensors(require_file(mert_dir / "model.safetensors")) + config = json.loads(require_file(adapter_dir / "config.json").read_text(encoding="utf-8")) + layers = int(config["backbone_config"]["num_hidden_layers"]) + scale = float(config["lora_alpha"]) / float(config["lora_rank"]) + merged: dict[str, torch.Tensor] = {name: tensor.clone() for name, tensor in mert_weights.items()} + for layer in range(layers): + for projection in PROJECTIONS: + base = f"layers.{layer}.attn.{projection}.weight" + a = adapter_weights[f"adapter.layers.{layer}.attn.{projection}.lora_A.weight"].float() + b = adapter_weights[f"adapter.layers.{layer}.attn.{projection}.lora_B.weight"].float() + merged[base] = merged[base].float().add((b @ a) * scale) + for name, tensor in adapter_weights.items(): + if not name.startswith("adapter."): + merged[name] = tensor + config["weights_format"] = "merged" + output_path.parent.mkdir(parents=True, exist_ok=True) + save_file(merged, output_path, metadata={"format": "pt"}) + (output_path.parent / "config.json").write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8") + + +def run(command: list[str]) -> None: + print("[run]", " ".join(command), flush=True) + subprocess.run(command, cwd=REPO_ROOT, check=True) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--adapter", + type=Path, + default=REPO_ROOT / "reference" / "SheetSage2", + help="SheetSage2 adapter snapshot containing config.json and model.safetensors.", + ) + parser.add_argument( + "--mert", + type=Path, + default=Path.home() / "Desktop" / "YuE2" / "MERT-v2-FullSong", + help="MERT-v2-FullSong parent snapshot containing model.safetensors.", + ) + parser.add_argument( + "--output", + type=Path, + default=Path.home() / "Desktop" / "SheetSage2-GGUF" / "sheetsage2-orig.gguf", + help="Output GGUF path.", + ) + parser.add_argument( + "--audiocpp-gguf", + type=Path, + default=REPO_ROOT / "build" / "debug" / "bin" / "audiocpp_gguf", + help="Path to audio.cpp GGUF converter.", + ) + parser.add_argument("--type", default="orig", choices=["orig", "f16", "bf16", "q4_0", "q4_k"]) + parser.add_argument("--keep-merged", type=Path, help="Optional path to keep the merged safetensors directory.") + parser.add_argument("--overwrite", action="store_true") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + adapter = args.adapter.resolve() + mert = args.mert.resolve() + output = args.output.resolve() + converter = require_file(args.audiocpp_gguf.resolve()) + spec = require_file(REPO_ROOT / "model_specs" / "sheetsage2.json") + temp_owner = tempfile.TemporaryDirectory(prefix="sheetsage2-merged-") + merged_dir = args.keep_merged.resolve() if args.keep_merged else Path(temp_owner.name) + merged_path = merged_dir / "model.safetensors" + merge_adapter(adapter, mert, merged_path) + root = output.parent + root.mkdir(parents=True, exist_ok=True) + shutil.copy2(merged_dir / "config.json", root / "config.json") + command = [ + str(converter), + "--input", + str(merged_path), + "--root", + str(root), + "--model-spec", + str(spec), + "--family", + "sheetsage2", + "--output", + str(output), + "--type", + args.type, + ] + if args.overwrite: + command.append("--overwrite") + run(command) + temp_owner.cleanup() + print("[done]", output) + + +if __name__ == "__main__": + main() diff --git a/tools/community_models/yue2/convert_yue2_gguf.py b/tools/community_models/yue2/convert_yue2_gguf.py new file mode 100644 index 00000000..f8b4e862 --- /dev/null +++ b/tools/community_models/yue2/convert_yue2_gguf.py @@ -0,0 +1,156 @@ +#!/usr/bin/env python3 +"""Convert YuE2 weights into the audio.cpp native GGUF package layout.""" + +from __future__ import annotations + +import argparse +import shutil +import subprocess +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[3] + + +SIDECARS = [ + ("YuE2-3B/config.json", "sidecars/yue2-model-config.json"), + ("YuE2-3B/yue2_generation_config.json", "sidecars/yue2-generation-config.json"), + ("YuE2-3B/qwen.tiktoken", "sidecars/yue2-qwen.tiktoken"), + ("YuE2-Vae/config.json", "sidecars/yue2-vae-config.json"), +] + + +def require_file(path: Path) -> Path: + if not path.is_file(): + raise FileNotFoundError(path) + return path + + +def copy_sidecars(source: Path, output: Path) -> None: + for src_rel, dst_rel in SIDECARS: + src = require_file(source / src_rel) + dst = output / dst_rel + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(src, dst) + + +def run(command: list[str]) -> None: + print("[run]", " ".join(command), flush=True) + subprocess.run(command, cwd=REPO_ROOT, check=True) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--source", + type=Path, + default=Path.home() / "Desktop" / "YuE2", + help="Directory containing YuE2-3B and YuE2-Vae snapshots.", + ) + parser.add_argument( + "--output", + type=Path, + default=Path.home() / "Desktop" / "Yue2-3B-GGUF", + help="Output GGUF package directory.", + ) + parser.add_argument( + "--audiocpp-gguf", + type=Path, + default=REPO_ROOT / "build" / "debug" / "bin" / "audiocpp_gguf", + help="Path to the audio.cpp GGUF converter.", + ) + parser.add_argument( + "--model-type", + default="q8_0", + choices=["orig", "f16", "bf16", "q8_0", "q2_k", "q3_k", "q4_k", "q5_k", "q6_k"], + help="Storage type for the YuE2 AR/NAR model GGUF.", + ) + parser.add_argument( + "--vae-type", + default="f16", + choices=["orig", "f16", "bf16", "q8_0"], + help="Storage type for the Oobleck VAE GGUF.", + ) + parser.add_argument("--overwrite", action="store_true") + return parser.parse_args() + + +def typed_name(stem: str, storage_type: str) -> str: + if storage_type == "orig": + if stem == "yue2-3b": + return f"{stem}-bf16.gguf" + if stem == "yue2-vae": + return f"{stem}-f32.gguf" + return f"{stem}-{storage_type}.gguf" + + +def main() -> None: + args = parse_args() + source = args.source.resolve() + output = args.output.resolve() + converter = require_file(args.audiocpp_gguf.resolve()) + model_spec = require_file(REPO_ROOT / "model_specs" / "yue2.json") + + model_weights = require_file(source / "YuE2-3B" / "model.safetensors") + vae_weights = require_file(source / "YuE2-Vae" / "model.safetensors") + output.mkdir(parents=True, exist_ok=True) + copy_sidecars(source, output) + model_output = output / typed_name("yue2-3b", args.model_type) + vae_output = output / typed_name("yue2-vae", args.vae_type) + + model_cmd = [ + str(converter), + "--input", + f"model_weights={model_weights}", + "--input", + f"vae_weights={vae_weights}", + "--root", + str(output), + "--model-spec", + str(model_spec), + "--family", + "yue2", + "--output", + str(model_output), + "--type", + args.model_type, + "--no-sidecars", + "--allow-missing-model-spec", + "--exclude-prefix", + "vae_weights/", + ] + vae_cmd = [ + str(converter), + "--input", + f"model_weights={model_weights}", + "--input", + f"vae_weights={vae_weights}", + "--root", + str(output), + "--model-spec", + str(model_spec), + "--family", + "yue2", + "--output", + str(vae_output), + "--type", + args.vae_type, + "--no-sidecars", + "--allow-missing-model-spec", + "--exclude-prefix", + "model_weights/", + "--fold-weight-norm", + "vae_weights/*", + ] + if args.overwrite: + model_cmd.append("--overwrite") + vae_cmd.append("--overwrite") + + run(model_cmd) + run(vae_cmd) + print("[done]", output) + print("[entry]", model_output) + + +if __name__ == "__main__": + main() diff --git a/webui/configs/model_params.json b/webui/configs/model_params.json index 8883d6e6..5d443310 100644 --- a/webui/configs/model_params.json +++ b/webui/configs/model_params.json @@ -8,7 +8,7 @@ {"name": "reference_duration_sec", "type": "slider", "label": "reference_duration_sec", "label_en": "Reference trim (s)", "default": 15.0, "minimum": 1.0, "maximum": 60.0, "step": 1.0, "info": "Trim the speaker reference before encoding. Around 10 s usually clones best."}, {"name": "seed", "type": "number", "label": "seed", "label_en": "Seed", "default": 0, "minimum": 0, "step": 1, "precision": 0} ], - "_comment": "WebUI TTS 高级参数控件配置：按模型 family 动态生成控件（gr.render）。每项字段：name=选项键(随请求 options 透传给模型)；type=slider|number|bool|text|choice；label/info=显示文案；default=默认值(应等于模型默认，已按 src/models//*.cpp 校对)；minimum/maximum/step=数值范围；precision=0 表示整数；choices=下拉候选。规则：只有被用户改动过的控件值才会随请求发送；seed/max_tokens 已有专用输入框，勿在此重复；参考文本用『参考文本』框(reference_text)；文件路径/parity 类参数(如 *_noise_file)未纳入，可用『其它参数(JSON)』兜底框传。", + "_comment": "WebUI TTS 高级参数控件配置：按模型 family 动态生成控件（gr.render）。每项字段：name=选项键；type=slider|number|bool|text|choice；scope=session 时写入 session_options，否则随请求 options 透传给模型；label/info=显示文案；default=默认值(应等于模型默认，已按 src/models//*.cpp 校对)；minimum/maximum/step=数值范围；precision=0 表示整数；choices=下拉候选。seed/max_tokens 已有专用输入框，勿在此重复；参考文本用『参考文本』框(reference_text)；文件路径/parity 类参数(如 *_noise_file)未纳入，可用『其它参数(JSON)』兜底框传。", "qwen3_tts": [ {"name": "temperature", "type": "slider", "label": "temperature", "default": 0.9, "minimum": 0.0, "maximum": 2.0, "step": 0.05}, @@ -194,6 +194,31 @@ {"name": "top_k", "type": "number", "label": "top_k", "default": 50, "minimum": 1, "maximum": 1024, "step": 1, "precision": 0} ], + "yue2": [ + {"name": "main_gguf", "type": "choice", "scope": "session", "session_option": "yue2.model_gguf", "label": "main_gguf", "label_en": "Main weights", "default": "yue2-3b-q8_0.gguf", "choices": ["yue2-3b-q8_0.gguf", "yue2-3b-q4_0.gguf", "yue2-3b-bf16.gguf"], "info": "Reload the model after changing this value."}, + {"name": "vae_gguf", "type": "choice", "scope": "session", "session_option": "yue2.vae_gguf", "label": "vae_gguf", "label_en": "VAE weights", "default": "yue2-vae-f16.gguf", "choices": ["yue2-vae-f16.gguf", "yue2-vae-f32.gguf"], "info": "Reload the model after changing this value."}, + {"name": "style", "type": "text", "label": "style", "label_en": "Style", "default": "English, indie pop, bright acoustic guitar, soft drums, warm lead vocal, polished demo mix", "placeholder": "English, city pop, groovy bass, synth, energetic vocal"}, + {"name": "abc", "type": "text", "label": "abc", "label_en": "ABC score", "default": "", "placeholder": "Optional ABC notation. Use cot=melody or cot=full.", "lines": 4}, + {"name": "abc_file", "type": "text", "label": "abc_file", "label_en": "ABC file path", "default": "", "placeholder": "/path/to/score.abc"}, + {"name": "cot", "type": "choice", "label": "cot", "label_en": "Planning route", "default": "off", "choices": ["off", "melody", "full"], "info": "off = direct generation; melody/full use or generate ABC planning."}, + {"name": "cfg_scale", "type": "slider", "label": "cfg_scale", "label_en": "Semantic guidance", "default": 1.01, "minimum": 0.0, "maximum": 5.0, "step": 0.01}, + {"name": "num_inference_steps", "type": "number", "label": "num_inference_steps", "label_en": "NAR steps", "default": 8, "minimum": 1, "maximum": 64, "step": 1, "precision": 0}, + {"name": "abc_temperature", "type": "slider", "label": "abc_temperature", "label_en": "ABC temperature", "default": 0.7, "minimum": 0.0, "maximum": 2.0, "step": 0.05}, + {"name": "abc_top_p", "type": "slider", "label": "abc_top_p", "label_en": "ABC top-p", "default": 0.9, "minimum": 0.01, "maximum": 1.0, "step": 0.01}, + {"name": "abc_top_k", "type": "number", "label": "abc_top_k", "label_en": "ABC top-k", "default": 30, "minimum": 1, "step": 1, "precision": 0}, + {"name": "abc_repetition_penalty", "type": "slider", "label": "abc_repetition_penalty", "label_en": "ABC repetition penalty", "default": 1.005, "minimum": 0.1, "maximum": 2.0, "step": 0.001}, + {"name": "abc_penalty_window", "type": "number", "label": "abc_penalty_window", "label_en": "ABC penalty window", "default": 100, "minimum": 1, "step": 1, "precision": 0}, + {"name": "abc_min_tokens", "type": "number", "label": "abc_min_tokens", "label_en": "ABC min tokens", "default": 32, "minimum": 0, "step": 1, "precision": 0}, + {"name": "abc_max_tokens", "type": "number", "label": "abc_max_tokens", "label_en": "ABC max tokens", "default": 4096, "minimum": 1, "step": 1, "precision": 0}, + {"name": "semantic_temperature", "type": "slider", "label": "semantic_temperature", "label_en": "Semantic temperature", "default": 1.0, "minimum": 0.0, "maximum": 2.0, "step": 0.05}, + {"name": "semantic_top_p", "type": "slider", "label": "semantic_top_p", "label_en": "Semantic top-p", "default": 0.95, "minimum": 0.01, "maximum": 1.0, "step": 0.01}, + {"name": "semantic_top_k", "type": "number", "label": "semantic_top_k", "label_en": "Semantic top-k", "default": 100, "minimum": 1, "step": 1, "precision": 0}, + {"name": "semantic_repetition_penalty", "type": "slider", "label": "semantic_repetition_penalty", "label_en": "Semantic repetition penalty", "default": 1.2, "minimum": 0.1, "maximum": 2.0, "step": 0.01}, + {"name": "semantic_penalty_window", "type": "number", "label": "semantic_penalty_window", "label_en": "Semantic penalty window", "default": 50, "minimum": 1, "step": 1, "precision": 0}, + {"name": "semantic_min_tokens", "type": "number", "label": "semantic_min_tokens", "label_en": "Semantic min tokens", "default": 200, "minimum": 0, "step": 1, "precision": 0}, + {"name": "semantic_max_tokens", "type": "number", "label": "semantic_max_tokens", "label_en": "Semantic max tokens", "default": 9000, "minimum": 1, "step": 1, "precision": 0} + ], + "minimax_h3": [ {"name": "num_inference_steps", "type": "number", "label": "Denoising steps", "default": 12, "minimum": 1, "maximum": 50, "step": 1, "precision": 0, "info": "Twelve denoising steps provide a practical quality and performance balance."}, {"name": "num_frames", "type": "number", "label": "Output frames", "default": 241, "minimum": 5, "maximum": 1441, "step": 4, "precision": 0, "info": "Approximately 24 frames per output second; 241 frames produces about 10 seconds of audio."}, diff --git a/webui/configs/models_catalog.json b/webui/configs/models_catalog.json index 2be709a6..eaa81a6f 100644 --- a/webui/configs/models_catalog.json +++ b/webui/configs/models_catalog.json @@ -89,6 +89,10 @@ { "id": "ace-step", "display_name": "ACE-Step 1.5 (music gen)", "family": "ace_step", "path": "models/Ace-Step1.5", "task": "gen", "mode": "offline", "download_id": "ace_step", "default_text": "upbeat pop music with bright vocals and energetic drums", "session_options": { "ace_step.mem_saver": "true", "ace_step.dit_weight_type": "q8_0", "ace_step.text_encoder_weight_type": "q8_0", "ace_step.planner_weight_type": "q8_0" }, "min_vram_gb": 8 }, { "id": "minimax-music3", "display_name": "MiniMax-Music3 (song gen)", "family": "minimax_music3", "path": "models/MiniMax-Music3-GGUF", "task": "gen", "mode": "offline", "download_id": "minimax_music3_q4_0", "min_vram_gb": 12 }, + { "id": "yue2", "display_name": "Yue2 3B (song gen)", "family": "yue2", "path": "models/Yue2-3B-GGUF", "task": "gen", "mode": "offline", "download_id": "yue2_main_q8_0", "default_text": "[Verse]\nSoft morning light is touching the window.\nI hear the city waking below.\n[Chorus]\nStay with the rhythm, let it carry us home.\nSing with the sunrise, we are never alone.", "min_vram_gb": 12, + "input_hint_en": "**Yue2 3B**: provide lyrics and a style prompt. The default inference combo is Main Q8_0 + VAE F16; choose other main/VAE files from Model parameters before loading." }, + { "id": "sheetsage2", "display_name": "SheetSage2 (audio to ABC)", "family": "sheetsage2", "path": "models/SheetSage2-GGUF/sheetsage2-orig.gguf", "task": "midi", "mode": "offline", "download_id": "sheetsage2_orig", "min_vram_gb": 8, + "input_hint_en": "**SheetSage2**: upload a song or instrumental recording to transcribe it into an ABC score artifact." }, { "id": "stable-audio-small-music","display_name": "Stable Audio 3 Small Music (gen)", "family": "stable_audio", "path": "models/stable-audio-3-small-music", "task": "gen", "mode": "offline", "download_id": "stable_audio_3_small_music", "min_vram_gb": 4 }, { "id": "stable-audio-small-sfx", "display_name": "Stable Audio 3 Small SFX (gen)", "family": "stable_audio", "path": "models/stable-audio-3-small-sfx", "task": "gen", "mode": "offline", "download_id": "stable_audio_3_small_sfx", "min_vram_gb": 4 }, { "id": "stable-audio-medium", "display_name": "Stable Audio 3 Medium (gen)", "family": "stable_audio", "path": "models/stable-audio-3-medium", "task": "gen", "mode": "offline", "download_id": "stable_audio_3_medium", "session_options": { "stable_audio.mem_saver": "true" }, "min_vram_gb": 10 }, diff --git a/webui/native/dist/index.html b/webui/native/dist/index.html index 6902014d..f0e788ef 100644 --- a/webui/native/dist/index.html +++ b/webui/native/dist/index.html @@ -24,27 +24,27 @@ })(); -
diff --git a/webui/native/package-lock.json b/webui/native/package-lock.json index 4a315649..7287976c 100644 --- a/webui/native/package-lock.json +++ b/webui/native/package-lock.json @@ -7,6 +7,9 @@ "": { "name": "audiocpp-native-webui", "version": "0.1.0", + "dependencies": { + "abcjs": "^6.7.0" + }, "devDependencies": { "@sveltejs/adapter-static": "^3.0.9", "@sveltejs/kit": "^2.37.0", @@ -1044,6 +1047,16 @@ "dev": true, "license": "MIT" }, + "node_modules/abcjs": { + "version": "6.7.0", + "resolved": "https://registry.npmjs.org/abcjs/-/abcjs-6.7.0.tgz", + "integrity": "sha512-dSj1Iho8IvGZEbfRTcGS4uGt+DMrrXTqXXpfNXYMuUNfDMRadTqA1DlVvs3u0uUDJIej8UOAeDauEPnZ6RThmg==", + "license": "MIT", + "funding": { + "type": "github", + "url": "https://github.com/sponsors/paulrosen" + } + }, "node_modules/acorn": { "version": "8.17.0", "resolved": "https://registry.npmjs.org/acorn/-/acorn-8.17.0.tgz", diff --git a/webui/native/package.json b/webui/native/package.json index 8fc8ff4f..967e48f4 100644 --- a/webui/native/package.json +++ b/webui/native/package.json @@ -16,5 +16,8 @@ "svelte-check": "^4.3.1", "typescript": "^5.9.2", "vite": "^7.1.3" + }, + "dependencies": { + "abcjs": "^6.7.0" } } diff --git a/webui/native/src/lib/catalog.ts b/webui/native/src/lib/catalog.ts index 894350e0..b97e3869 100644 --- a/webui/native/src/lib/catalog.ts +++ b/webui/native/src/lib/catalog.ts @@ -161,6 +161,14 @@ function packageLabel(entry: PackageEntry): string { return `GGUF ${precision}`; } if (entry.family === 'irodori_tts' && entry.id.includes('_anime_')) return 'Anime Q8'; + if (entry.family === 'yue2') { + if (entry.id === 'yue2_main_q8_0') return 'Main Q8_0'; + if (entry.id === 'yue2_main_q4_0') return 'Main Q4_0'; + if (entry.id === 'yue2_main_bf16') return 'Main BF16'; + if (entry.id === 'yue2_vae_f16') return 'VAE F16'; + if (entry.id === 'yue2_vae_f32') return 'VAE F32'; + return entry.display_name || 'Yue2 component'; + } if (entry.format === 'safetensors') return 'Safetensors'; if (entry.id.includes('int8_dit')) return 'GGUF Q4 ConvRot'; if (entry.precision === 'q4_k' || entry.precision === 'q4_0') return 'GGUF Q4'; @@ -175,7 +183,7 @@ function packageModelPath(entry: PackageEntry): string { if (entry.format === 'gguf' && entry.family === 'minimax_h3') { const entryName = entry.id.includes('int8_dit') ? 'dit_int8.gguf' : 'dit.gguf'; modelFile = entry.files?.find((file) => file.toLowerCase().endsWith(`/${entryName}`)); - } else if (entry.format === 'gguf' && entry.family === 'minimax_music3') { + } else if (entry.format === 'gguf' && (entry.family === 'minimax_music3' || entry.family === 'yue2')) { return `models/${entry.target_directory}`; } else if (entry.format === 'gguf') { modelFile = entry.files?.find((file) => file.toLowerCase().endsWith('.gguf')); @@ -190,33 +198,56 @@ function packageModelPath(entry: PackageEntry): string { } function packageSessionOptions(entry: PackageEntry): Record | undefined { - if (entry.family !== 'minimax_music3') return undefined; - if (entry.id === 'minimax_music3_q8_0') { - return { - 'minimax_music3.language_model_gguf': 'language_model_q8_0.gguf', - 'minimax_music3.rvq_depth_decoder_gguf': 'rvq_depth_decoder_q8_0.gguf', - 'minimax_music3.flow_transformer_gguf': 'transformer_q8_0.gguf' - }; - } - if (entry.id === 'minimax_music3_bf16') { - return { - 'minimax_music3.language_model_gguf': 'language_model_bf16.gguf', - 'minimax_music3.rvq_depth_decoder_gguf': 'rvq_depth_decoder_bf16.gguf', - 'minimax_music3.flow_transformer_gguf': 'transformer_bf16.gguf' - }; - } - if (entry.id === 'minimax_music3_q4_0') { - return { - 'minimax_music3.language_model_gguf': 'language_model_q4_0.gguf', - 'minimax_music3.rvq_depth_decoder_gguf': 'rvq_depth_decoder_q8_0.gguf', - 'minimax_music3.flow_transformer_gguf': 'transformer_q4_0.gguf' - }; + if (entry.family === 'minimax_music3') { + if (entry.id === 'minimax_music3_q8_0') { + return { + 'minimax_music3.language_model_gguf': 'language_model_q8_0.gguf', + 'minimax_music3.rvq_depth_decoder_gguf': 'rvq_depth_decoder_q8_0.gguf', + 'minimax_music3.flow_transformer_gguf': 'transformer_q8_0.gguf' + }; + } + if (entry.id === 'minimax_music3_bf16') { + return { + 'minimax_music3.language_model_gguf': 'language_model_bf16.gguf', + 'minimax_music3.rvq_depth_decoder_gguf': 'rvq_depth_decoder_bf16.gguf', + 'minimax_music3.flow_transformer_gguf': 'transformer_bf16.gguf' + }; + } + if (entry.id === 'minimax_music3_q4_0') { + return { + 'minimax_music3.language_model_gguf': 'language_model_q4_0.gguf', + 'minimax_music3.rvq_depth_decoder_gguf': 'rvq_depth_decoder_q8_0.gguf', + 'minimax_music3.flow_transformer_gguf': 'transformer_q4_0.gguf' + }; + } } return undefined; } function installChoices(entry: CatalogEntry): InstallPackageChoice[] { const exposesAllGguf = exposeAllGgufPackageFamilies.has(entry.family); + if (entry.family === 'yue2') { + const related = packages.filter((candidate) => + candidate.family === entry.family && candidate.format === 'gguf'); + const order = new Map([ + ['yue2_main_q8_0', 0], + ['yue2_main_q4_0', 1], + ['yue2_main_bf16', 2], + ['yue2_vae_f16', 3], + ['yue2_vae_f32', 4] + ]); + return related + .filter((candidate) => candidate.format === 'gguf') + .sort((left, right) => (order.get(left.id) ?? 99) - (order.get(right.id) ?? 99)) + .map((candidate) => ({ + id: candidate.id, + label: packageLabel(candidate), + path: packageModelPath(candidate), + format: candidate.format, + precision: candidate.precision, + session_options: packageSessionOptions(candidate) + })); + } const related = exposesAllGguf ? relatedExposeAllGgufPackages(entry) : relatedPackages(entry); if (entry.family === 'ace_step' || entry.family === 'minimax_music3' || exposesAllGguf) { diff --git a/webui/native/src/lib/models/panels.ts b/webui/native/src/lib/models/panels.ts new file mode 100644 index 00000000..cc5c238b --- /dev/null +++ b/webui/native/src/lib/models/panels.ts @@ -0,0 +1,34 @@ +import Yue2Panel from './yue2/Yue2Panel.svelte'; + +export interface GenericControlReplacements { + packageButtons?: boolean; + text?: boolean; + genSource?: boolean; + language?: boolean; + seed?: boolean; + duration?: boolean; + params?: boolean; + advancedJson?: boolean; +} + +export const modelStudioPanels = { + yue2: { + component: Yue2Panel, + requestMode: 'yue2', + replacesGenericControls: { + packageButtons: true, + text: true, + genSource: true, + language: true, + seed: true, + duration: true, + params: true, + advancedJson: true + } + } +}; + +export function modelStudioPanelFor(family?: string) { + if (!family) return undefined; + return modelStudioPanels[family as keyof typeof modelStudioPanels]; +} diff --git a/webui/native/src/lib/models/yue2/Yue2Panel.svelte b/webui/native/src/lib/models/yue2/Yue2Panel.svelte new file mode 100644 index 00000000..e5adf127 --- /dev/null +++ b/webui/native/src/lib/models/yue2/Yue2Panel.svelte @@ -0,0 +1,642 @@ + + +
+
+ + +
+ +
+
+ + +
+ {#each componentSpecs as spec} +
+ + + {#if localizedParameterText(spec, 'info', tr)}{localizedParameterText(spec, 'info', tr)}{/if} +
+ {/each} +
+ +
+ {#each coreSpecs as spec} +
+ + {#if spec.type === 'choice'} + + {:else if spec.type === 'slider'} +
+ setParameterValue(spec, event.currentTarget.valueAsNumber)} /> + {String(advancedValues[spec.name])} +
+ {:else} + setParameterValue(spec, + spec.type === 'number' ? event.currentTarget.valueAsNumber : event.currentTarget.value)} /> + {/if} + {#if localizedParameterText(spec, 'info', tr)}{localizedParameterText(spec, 'info', tr)}{/if} +
+ {/each} +
+ +
+ Yue2 ABC conditioning {abcSpecs.length} +
+
+
+
+ Cover source + Use SheetSage2 + MERT2 to extract an editable ABC score from a song. +
+ +
+ + Warning: VRAM may remain in use after unloading. + coverAudioFile = event.currentTarget.files?.[0] || null} /> + +
+ + {#if coverStatus}{coverStatus}{/if} + {#if coverError}{coverError}{/if} +
+ +
+ +
+
+
+ ABC score editor + Edit the extracted score here. The sheet preview updates from this ABC. +
+
+ + {#if specByName('abc_file')} +
+ + setNamedParameter('abc_file', event.currentTarget.value)} /> +
+ {/if} +
+ +
+
+
+ Sheet preview + Rendered from the editable ABC score. +
+
+
+ {#if !abcDraft.trim()}Extract or paste ABC to preview the score.{/if} +
+ {#if abcRenderError}{abcRenderError}{/if} +
+
+
+ +
+ Yue2 semantic sampling {semanticSpecs.length} +
+ {#each semanticSpecs as spec} +
+ + {#if spec.type === 'slider'} +
+ setParameterValue(spec, event.currentTarget.valueAsNumber)} /> + {String(advancedValues[spec.name])} +
+ {:else} + setParameterValue(spec, event.currentTarget.valueAsNumber)} /> + {/if} +
+ {/each} +
+
+ +
+ Yue2 ABC planner sampling {plannerSpecs.length} +
+ {#each plannerSpecs as spec} +
+ + {#if spec.type === 'slider'} +
+ setParameterValue(spec, event.currentTarget.valueAsNumber)} /> + {String(advancedValues[spec.name])} +
+ {:else} + setParameterValue(spec, event.currentTarget.valueAsNumber)} /> + {/if} +
+ {/each} +
+
+
+ + diff --git a/webui/native/src/lib/types.ts b/webui/native/src/lib/types.ts index caeceefe..32f1ff2c 100644 --- a/webui/native/src/lib/types.ts +++ b/webui/native/src/lib/types.ts @@ -35,6 +35,8 @@ export interface CatalogEntry { export interface ParamSpec { name: string; type: 'slider' | 'number' | 'bool' | 'text' | 'choice'; + scope?: 'request' | 'session'; + session_option?: string; label: string; label_en?: string; info?: string; diff --git a/webui/native/src/routes/+page.svelte b/webui/native/src/routes/+page.svelte index 58622d53..892f23b4 100644 --- a/webui/native/src/routes/+page.svelte +++ b/webui/native/src/routes/+page.svelte @@ -34,6 +34,7 @@ import MediaPreview from '$lib/MediaPreview.svelte'; import { defaultChunkBudget, splitTtsChunks } from '$lib/text'; import { UI_THEME_STORAGE_KEY, resolvedTheme, resolveUiTheme, uiThemes, type UiTheme } from '$lib/theme'; + import { modelStudioPanelFor, type GenericControlReplacements } from '$lib/models/panels'; import Arena from './Arena.svelte'; import type { AudioOutput, @@ -154,6 +155,7 @@ 'meanvc2', 'midashenglm_gen' ]); + const noGenericControlReplacements: GenericControlReplacements = {}; function chooseUiLanguage(code: string) { uiLanguage = resolveUiLanguage([code]); @@ -250,6 +252,11 @@ } } + function ensureYue2DefaultLyrics(entry = selected) { + if (entry?.family !== 'yue2' || lyrics.trim()) return; + lyrics = entry.default_text || ''; + } + function requestText() { return text.trim() ? text : (selected.default_text || ''); } @@ -398,6 +405,10 @@ })).filter((group) => group.entries.length > 0); $: isLoaded = loadedModels.some((model) => model.id === selectedId && model.loaded && modelMatchesSelectedPackage(model, selected)); + $: modelStudioPanelConfig = modelStudioPanelFor(selected?.family); + $: modelStudioPanel = modelStudioPanelConfig?.component; + $: replacesGenericControls = modelStudioPanelConfig?.replacesGenericControls || noGenericControlReplacements; + $: usesYue2Request = modelStudioPanelConfig?.requestMode === 'yue2'; $: isFireRedAudioEdit = selected?.id === 'firered-audio-semantic-edit' || selected?.id === 'firered-audio-acoustic-edit'; $: allowsAutoDuration = selected?.family === 'ace_step'; @@ -410,7 +421,7 @@ ) && selected?.task === 'tts'; $: needsSource = ['asr', 'vc', 'svc', 's2s', 'sep', 'vad', 'diar', 'align', 'midi'].includes(selected?.task) || isFireRedAudioEdit; - $: acceptsSource = needsSource || selected?.task === 'gen'; + $: acceptsSource = needsSource || (selected?.task === 'gen' && !replacesGenericControls.genSource); $: acceptsVideo = selected?.request_options?.includes('video') === true; $: needsVoice = (['clon', 'vc', 'svc'].includes(selected?.task) && selected?.family !== 'rvc') || (selected?.task === 's2s' && selected?.family === 'personaplex') || @@ -438,7 +449,8 @@ $: quickStartVoicePreview = quickStartVoice && server?.ui_management !== false && !usesBuiltInVoiceSelector ? voicePreviewUrl(demoVoiceSources[quickStartVoice] || quickStartVoice) : ''; - $: showsText = ['tts', 'clon', 'gen', 's2s', 'align', 'vdes'].includes(selected?.task); + $: showsText = ['tts', 'clon', 'gen', 's2s', 'align', 'vdes'].includes(selected?.task) && + !replacesGenericControls.text; $: supportsLiveAsr = selected?.task === 'asr' && ['voxtral_realtime', 'nemotron_asr', 'higgs_audio_stt', 'sense_asr', 'vibevoice_asr_streaming'].includes(selected?.family); $: modelInventoryLoading = server === null || @@ -498,13 +510,19 @@ function mergedSessionOptions(entry: CatalogEntry) { const packageChoice = selectedPackageChoice(entry); - return { ...(entry.session_options || {}), ...(packageChoice?.session_options || {}) }; + const sessionParams = entry.id === selectedId ? sessionParameterOptions() : {}; + return { ...(entry.session_options || {}), ...(packageChoice?.session_options || {}), ...sessionParams }; } function packageSessionOptionsMatch(entry: CatalogEntry, choice: InstallPackageChoice, model: LoadedModel) { - const expected = choice.session_options || {}; + const expected = mergedSessionOptions(entry); const keys = Array.from(new Set((entry.install_packages || []) .flatMap((candidate) => Object.keys(candidate.session_options || {})))); + if (entry.id === selectedId) { + for (const spec of paramSpecs.filter((candidate) => candidate.scope === 'session')) { + keys.push(spec.session_option || spec.name); + } + } if (!keys.length) return true; const actual = model.session_options || {}; return keys.every((key) => actual[key] === expected[key]); @@ -1010,7 +1028,11 @@ } else if (selected?.task === 'gen') { duration = 30; } - if (!text.trim() && selected?.default_text) { + if (usesYue2Request) { + text = ''; + lyrics = ''; + ensureYue2DefaultLyrics(); + } else if (!text.trim() && selected?.default_text) { text = selected.default_text; } if (selected?.builtin_voices?.length && selected.default_voice && !quickStartVoice) { @@ -1141,7 +1163,8 @@ try { const targetPath = comparablePath(modelPath); const replaced = loadedModels.filter((model) => model.loaded && - (model.id !== selected.id || comparablePath(model.path) !== targetPath)); + (model.id !== selected.id || comparablePath(model.path) !== targetPath || + !modelMatchesSelectedPackage(model, selected))); for (const model of replaced) { log(`Unloading ${loadedModelName(model)} before loading ${selected.display_name}.`); await unloadModel(model.id); @@ -1248,7 +1271,21 @@ throw new Error(`Advanced JSON is invalid: ${error instanceof Error ? error.message : error}`); } const defaults = selected.default_options || {}; - return { ...defaults, ...advancedValues, ...raw }; + const requestValues = Object.fromEntries(Object.entries(advancedValues) + .filter(([name, value]) => { + const spec = paramSpecs.find((candidate) => candidate.name === name); + if (spec?.scope === 'session') return false; + if (usesYue2Request && typeof value === 'string' && value.trim().length === 0) return false; + return true; + })); + return { ...defaults, ...requestValues, ...raw }; + } + + function sessionParameterOptions() { + return Object.fromEntries(paramSpecs + .filter((spec) => spec.scope === 'session') + .map((spec) => [spec.session_option || spec.name, String(advancedValues[spec.name] ?? spec.default ?? '')]) + .filter(([, value]) => value.length > 0)); } function base64Text(value: string): string { @@ -1675,15 +1712,19 @@ } else { if (needsSource && !audio) throw new StatusWarning('Choose a source audio file.'); const request: Record = { options }; - if (['gen', 's2s', 'align'].includes(selected.task) && text.trim()) request.text = text; - if (['gen', 's2s', 'align'].includes(selected.task) && language.trim()) request.language = language; + if (['gen', 's2s', 'align'].includes(selected.task) && text.trim() && !usesYue2Request) request.text = text; + if (['gen', 's2s', 'align'].includes(selected.task) && language.trim() && !usesYue2Request) request.language = language; if (selected.task === 'gen') { - const resolvedText = requestText(); - if (resolvedText) request.text = resolvedText; - if (lyrics.trim()) request.lyrics = lyrics; - if (!isFireRedAudioEdit) { - if (usesDurationSecOption) options.duration_sec = duration; - else request.duration_seconds = duration; + if (usesYue2Request) { + request.lyrics = lyrics.trim(); + } else { + const resolvedText = requestText(); + if (resolvedText) request.text = resolvedText; + if (lyrics.trim()) request.lyrics = lyrics; + if (!isFireRedAudioEdit) { + if (usesDurationSecOption) options.duration_sec = duration; + else request.duration_seconds = duration; + } } request.seed = resolvedSeed; if (supportsMaxTokens(selected)) request.max_tokens = maxTokens; @@ -2159,7 +2200,7 @@ {selectedId ? tr('studio.estimatedVram', { value: selected?.min_vram_gb || '?' }) : tr('studio.vram')} - {#if selectedId && (selected.install_packages || []).length} + {#if selectedId && (selected.install_packages || []).length && !replacesGenericControls.packageButtons}
{#each studioPackageSlots(selected) as slot} {@const choice = slot.choice} @@ -2218,9 +2259,28 @@ {/if} {#if selected.task === 'gen'} - - + {#if modelStudioPanel} + + {:else} + + + {/if} {#if selected.family === 'ace_step'}