Merge branch 'main' into vibeasr-vae-encoder

Author: XsquirrelC

CMakeLists.txt

@@ -2118,6 +2118,13 @@ if (ENGINE_BUILD_TESTS)
         COMMAND chinese_normalization_test
     )
 
+    add_engine_unittest(text_chunking_test tests/unittests/test_text_chunking.cpp)
+
+    add_test(
+        NAME text_chunking_test
+        COMMAND text_chunking_test
+    )
+
     add_engine_unittest(unicode_normalization_test tests/unittests/test_unicode_normalization.cpp)
 
     add_test(

src/framework/text/chunking.cpp

@@ -63,6 +63,17 @@ bool is_clause_break(std::string_view token) {
            token == u8"，" || token == u8"、" || token == u8"；" || token == u8"：";
 }
 
+bool is_cjk_punctuation_delimiter(std::string_view token) noexcept {
+    // Full-width CJK sentence/clause punctuation (。！？，、；：) separates words
+    // even without ASCII spaces — CJK text has no inter-word spaces, so without
+    // this a whole paragraph parses as one word and Default-mode chunking can
+    // never split it at the codepoint budget (only TagAware/Japanese could).
+    // ASCII punctuation stays attached to its word so space-delimited Latin
+    // text is unaffected.
+    const auto leading = static_cast<unsigned char>(token.front());
+    return leading >= 0x80 && (is_sentence_break(token) || is_clause_break(token));
+}
+
 bool is_tag_open(std::string_view token) {
     return token == "[" || token == "<";
 }
@@ -112,9 +123,25 @@ std::vector<WordRange> split_word_ranges(const std::vector<Utf8Span> & spans) {
         if (span_pos >= spans.size()) {
             break;
         }
+        // A CJK punctuation token forms a word of its own (attached to nothing),
+        // so a subsequent run never absorbs it and boundaries can land on it.
+        if (is_cjk_punctuation_delimiter(spans[span_pos].text)) {
+            words.push_back({
+                span_pos,
+                span_pos + 1,
+                spans[span_pos].start,
+                spans[span_pos].end,
+                is_sentence_break(spans[span_pos].text),
+                is_clause_break(spans[span_pos].text),
+            });
+            span_pos += 1;
+            continue;
+        }
         const size_t word_start = span_pos;
         size_t word_end = span_pos + 1;
-        while (word_end < spans.size() && !is_ascii_space(spans[word_end].text)) {
+        while (word_end < spans.size() &&
+               !is_ascii_space(spans[word_end].text) &&
+               !is_cjk_punctuation_delimiter(spans[word_end].text)) {
             ++word_end;
         }
         const auto last = spans[word_end - 1].text;

src/models/higgs_audio_tts/codec.cpp

@@ -62,7 +62,7 @@ constexpr int64_t kCodecDecodeCapacityBucketFrames = 32;
 constexpr int64_t kCodecHopLength = 960;
 constexpr int64_t kCodecPadSamples = kCodecHopLength / 2;
 constexpr int64_t kCodecDecodeWindowFrames = 128;
-constexpr int64_t kCodecDecodeOverlapFrames = 8;
+constexpr int64_t kCodecDecodeContextFrames = 32;
 constexpr int64_t kCodecFullDecodeMaxFrames = 512;
 constexpr int64_t kResidualDilations[] = {1, 3, 9};
 
@@ -1584,11 +1584,14 @@ HiggsCodecDecodeOutput HiggsCodecRuntime::decode_codes(const std::vector<int32_t
     std::vector<int32_t> window_codes;
     int64_t emitted_frames = 0;
     while (emitted_frames < frames) {
-        const int64_t window_begin =
-            std::max<int64_t>(0, emitted_frames - kCodecDecodeOverlapFrames);
+        const int64_t emit_begin = emitted_frames;
         const int64_t emit_end =
-            std::min<int64_t>(frames, emitted_frames + kCodecDecodeWindowFrames);
-        const int64_t window_frames = emit_end - window_begin;
+            std::min<int64_t>(frames, emit_begin + kCodecDecodeWindowFrames);
+        const int64_t window_begin =
+            std::max<int64_t>(0, emit_begin - kCodecDecodeContextFrames);
+        const int64_t window_end =
+            std::min<int64_t>(frames, emit_end + kCodecDecodeContextFrames);
+        const int64_t window_frames = window_end - window_begin;
         window_codes.resize(static_cast<size_t>(window_frames * kCodecCodebooks));
         for (int64_t frame = 0; frame < window_frames; ++frame) {
             const auto src =
@@ -1602,9 +1605,9 @@ HiggsCodecDecodeOutput HiggsCodecRuntime::decode_codes(const std::vector<int32_t
         const auto window = run_window(
             window_codes,
             window_frames,
-            kCodecDecodeWindowFrames + kCodecDecodeOverlapFrames);
-        const int64_t trim_frames = emitted_frames - window_begin;
-        const int64_t emit_frames = emit_end - emitted_frames;
+            kCodecDecodeWindowFrames + 2 * kCodecDecodeContextFrames);
+        const int64_t trim_frames = emit_begin - window_begin;
+        const int64_t emit_frames = emit_end - emit_begin;
         const int64_t sample_begin = trim_frames * kCodecHopLength;
         const int64_t sample_count = emit_frames * kCodecHopLength;
         if (sample_begin < 0 || sample_count <= 0 ||

src/models/vevo2/fm.cpp

@@ -696,6 +696,7 @@ struct Vevo2FMGraph {
     }
 
     ~Vevo2FMGraph() {
+        engine::core::release_backend_graph_resources(backend, graph, true);
         if (gallocr != nullptr) {
             ggml_gallocr_free(gallocr);
             gallocr = nullptr;
@@ -757,39 +758,45 @@ struct Vevo2FMStepGraph {
         if (config.hidden_size % config.num_heads != 0) {
             throw std::runtime_error("Vevo2 FM hidden_size must be divisible by num_heads");
         }
+        ggml_init_params input_params{ggml_tensor_overhead() * 64, nullptr, true};
+        input_ctx.reset(ggml_init(input_params));
+        if (input_ctx == nullptr) {
+            throw std::runtime_error("failed to initialize Vevo2 FM step input context");
+        }
         ggml_init_params params{graph_context_bytes, nullptr, true};
         ctx.reset(ggml_init(params));
         if (ctx == nullptr) {
             throw std::runtime_error("failed to initialize Vevo2 FM step graph context");
         }
 
+        engine::core::ModuleBuildContext input_build_ctx{input_ctx.get(), "vevo2.fm.step.input", backend_type};
         engine::core::ModuleBuildContext build_ctx{ctx.get(), "vevo2.fm.step", backend_type};
         prompt_input = engine::core::make_tensor(
-            build_ctx,
+            input_build_ctx,
             GGML_TYPE_F32,
             engine::core::TensorShape::from_dims({1, prompt_frames, config.mel_dim})).tensor;
         xt_input = engine::core::make_tensor(
-            build_ctx,
+            input_build_ctx,
             GGML_TYPE_F32,
             engine::core::TensorShape::from_dims({1, target_frames, config.mel_dim})).tensor;
         cond_input = engine::core::make_tensor(
-            build_ctx,
+            input_build_ctx,
             GGML_TYPE_F32,
             engine::core::TensorShape::from_dims({1, cond_frames, config.hidden_size})).tensor;
         uncond_cond_input = engine::core::make_tensor(
-            build_ctx,
+            input_build_ctx,
             GGML_TYPE_F32,
             engine::core::TensorShape::from_dims({1, target_frames, config.hidden_size})).tensor;
         timestep_input = engine::core::make_tensor(
-            build_ctx,
+            input_build_ctx,
             GGML_TYPE_F32,
             engine::core::TensorShape::from_dims({1, config.hidden_size})).tensor;
         cond_position_input = engine::core::make_tensor(
-            build_ctx,
+            input_build_ctx,
             GGML_TYPE_I32,
             engine::core::TensorShape::from_dims({cond_frames})).tensor;
         uncond_position_input = engine::core::make_tensor(
-            build_ctx,
+            input_build_ctx,
             GGML_TYPE_I32,
             engine::core::TensorShape::from_dims({target_frames})).tensor;
         ggml_set_input(prompt_input);
@@ -878,8 +885,11 @@ struct Vevo2FMStepGraph {
 
         graph = ggml_new_graph_custom(ctx.get(), 524288, false);
         ggml_build_forward_expand(graph, output);
-        buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), backend);
-        if (buffer == nullptr) {
+        input_buffer = ggml_backend_alloc_ctx_tensors(input_ctx.get(), backend);
+        gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(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 Vevo2 FM step graph");
         }
 
@@ -904,9 +914,14 @@ struct Vevo2FMStepGraph {
     }
 
     ~Vevo2FMStepGraph() {
-        if (buffer != nullptr) {
-            ggml_backend_buffer_free(buffer);
-            buffer = nullptr;
+        engine::core::release_backend_graph_resources(backend, graph, true);
+        if (gallocr != nullptr) {
+            ggml_gallocr_free(gallocr);
+            gallocr = nullptr;
+        }
+        if (input_buffer != nullptr) {
+            ggml_backend_buffer_free(input_buffer);
+            input_buffer = nullptr;
         }
     }
 
@@ -966,6 +981,7 @@ struct Vevo2FMStepGraph {
     int64_t cond_frames = 0;
     int64_t prompt_frames = 0;
     int64_t target_frames = 0;
+    std::unique_ptr<ggml_context, GgmlContextDeleter> input_ctx;
     std::unique_ptr<ggml_context, GgmlContextDeleter> ctx;
     ggml_tensor * prompt_input = nullptr;
     ggml_tensor * xt_input = nullptr;
@@ -976,7 +992,8 @@ struct Vevo2FMStepGraph {
     ggml_tensor * uncond_position_input = nullptr;
     ggml_tensor * output = nullptr;
     ggml_cgraph * graph = nullptr;
-    ggml_backend_buffer_t buffer = nullptr;
+    ggml_backend_buffer_t input_buffer = nullptr;
+    ggml_gallocr_t gallocr = nullptr;
 };
 
 Vevo2FlowMatchingRuntime::Vevo2FlowMatchingRuntime(
@@ -1064,6 +1081,8 @@ Vevo2MelSequence Vevo2FlowMatchingRuntime::generate_mel(
     const auto cond_run_start = Clock::now();
     const auto diffusion_cond = graph_->run_conditioning(diffusion_tokens, config_);
     const double cond_run_ms = engine::debug::elapsed_ms(cond_run_start);
+    graph_.reset();
+    engine::core::trim_backend_pools(execution_context_.backend());
     const int64_t cond_frames = static_cast<int64_t>(diffusion_cond.size()) / config_.hidden_size;
     if (cond_frames * config_.hidden_size != static_cast<int64_t>(diffusion_cond.size())) {
         throw std::runtime_error("Vevo2 FM conditioning output shape mismatch");
@@ -1123,6 +1142,8 @@ Vevo2MelSequence Vevo2FlowMatchingRuntime::generate_mel(
     const auto read_start = Clock::now();
     out.values = step_graph_->read_output(config_);
     const double final_read_ms = engine::debug::elapsed_ms(read_start);
+    step_graph_.reset();
+    engine::core::trim_backend_pools(execution_context_.backend());
     engine::debug::timing_log_scalar("vevo2.fm.timbre_mel_ms", timbre_mel_ms);
     engine::debug::timing_log_scalar("vevo2.fm.cond.graph.build_ms", cond_graph_build_ms);
     engine::debug::timing_log_scalar("vevo2.fm.cond_run_ms", cond_run_ms);

tests/unittests/test_text_chunking.cpp

@@ -0,0 +1,85 @@
+// Default-mode text chunking regression: CJK text (no ASCII spaces) must still
+// split at the codepoint budget, with full-width punctuation as word boundaries.
+#include "engine/framework/text/chunking.h"
+
+#include <cstdlib>
+#include <iostream>
+#include <string>
+#include <vector>
+
+namespace {
+
+void require(bool ok, const std::string & what) {
+    if (!ok) {
+        std::cerr << "FAIL: " << what << "\n";
+        std::exit(1);
+    }
+}
+
+void check_chunks(const std::string & text, int64_t budget, size_t min_chunks, const std::string & label) {
+    const auto chunks = engine::text::split_text_chunks(text, budget);
+    if (chunks.size() < min_chunks) {
+        std::cerr << "FAIL: " << label << " expected >= " << min_chunks << " chunks, got " << chunks.size()
+                  << "\n";
+        std::exit(1);
+    }
+    std::string joined;
+    for (const auto & c : chunks) {
+        joined += c;
+    }
+    require(joined == text, label + ": chunks must concatenate back to the trimmed input verbatim");
+    for (const auto & c : chunks) {
+        require(!c.empty(), label + ": empty chunk");
+    }
+}
+
+}  // namespace
+
+int main() {
+    // Regression: a 236-codepoint CJK paragraph with the 200-codepoint default
+    // budget used to come back as a single chunk, because split_word_ranges only
+    // cut at ASCII spaces and the whole paragraph parsed as one word. Must now
+    // split into at least two sentence-aligned chunks.
+    const std::string cjk_long =
+        "大家好，我是零一B语音合成模型，现在进行流式输出测试。这段文字比较长，目的是看模型能否把整段话分成多个音频块，边合成边推送。"
+        "今天的天气很好，适合外出散步，湖边的柳树已经抽出了新芽，水面倒映着蓝天白云。如果流式工作正常，客户端应该能很快收到第一个音频块。"
+        "第二段继续测试分块是否稳定，这里再补充一些内容，让总字数超过上限，这样模型必须把文字切开分多次合成。流式的意义在于长文本不必等全部算完。"
+        "最后再来一句收尾的话，确认整个流程完整结束，谢谢大家。";
+    check_chunks(cjk_long, 200, 2, "CJK long text over budget splits");
+
+    // Under budget: still one chunk, verbatim.
+    check_chunks("短文本，不需要分块。", 200, 1, "short CJK text stays whole");
+
+    // Clause punctuation is a valid rollback boundary when the budget cuts a
+    // run of clauses: budget 12 lands inside the second clause run, so the
+    // first chunk must end at the 、 (clause) boundary, not mid-run.
+    {
+        const std::string text = "一二三四五六七八九十，一二三四五六七八九十。";
+        const auto chunks = engine::text::split_text_chunks(text, 12);
+        require(chunks.size() == 2, "clause rollback produces two chunks");
+        require(chunks[0] == "一二三四五六七八九十，", "first chunk ends at clause punctuation");
+        require(chunks[1] == "一二三四五六七八九十。", "second chunk keeps the sentence");
+    }
+
+    // Latin text is unaffected: ASCII punctuation stays attached to its word
+    // and boundaries still fall on ASCII-space words / sentence breaks.
+    {
+        const std::string text = "Hello world. How are you today?";
+        const auto chunks = engine::text::split_text_chunks(text, 20);
+        require(chunks.size() == 2, "latin splits at the sentence break");
+        require(chunks[0] == "Hello world.", "first latin chunk ends at the period");
+        require(chunks[1] == "How are you today?", "second latin chunk verbatim");
+    }
+
+    // A single oversized word without any boundary still passes through whole
+    // (pre-existing behavior: no hard character-level cut).
+    {
+        const std::string word(300, 'a');
+        const auto chunks = engine::text::split_text_chunks(word, 200);
+        require(chunks.size() == 1, "unbreakable oversized word stays one chunk");
+        require(chunks[0] == word, "oversized word verbatim");
+    }
+
+    std::cout << "text_chunking_test: all passed\n";
+    return 0;
+}