Improve Chatterbox memory usage

Improve Chatterbox memory usage

Author: 0xShug0

app/server/main.cpp

@@ -2,6 +2,8 @@
 #include "http.h"
 #include "runtime.h"
 
+#include "engine/framework/debug/trace.h"
+
 #include <filesystem>
 #include <iostream>
 #include <optional>
@@ -31,6 +33,7 @@ bool has_arg(int argc, char ** argv, const std::string & name) {
 void print_help() {
     std::cout
         << "audiocpp_server --config <server.json> [--host <ip>] [--port <port>] [--device <id>] [--threads <n>]\n"
+        << "                [--log] [--log-file <path>]\n"
         << "\n"
         << "Endpoints:\n"
         << "  GET  /health\n"
@@ -52,6 +55,11 @@ int main(int argc, char ** argv) {
         if (!config_path.has_value()) {
             throw std::runtime_error("missing required --config argument");
         }
+        const auto log_file = arg_value(argc, argv, "--log-file");
+        engine::debug::configure_logging(engine::debug::LoggingConfig{
+            has_arg(argc, argv, "--log") || log_file.has_value(),
+            log_file,
+        });
 
         auto config = minitts::server::load_server_config(*config_path);
         if (const auto host = arg_value(argc, argv, "--host")) {

app/server/runtime.cpp

@@ -6,17 +6,21 @@
 #include "engine/framework/runtime/registry.h"
 
 #include <algorithm>
+#include <chrono>
 #include <cmath>
 #include <cstdint>
 #include <sstream>
 #include <stdexcept>
+#include <unordered_map>
 #include <utility>
 
 namespace minitts::server {
 namespace {
 
 using engine::io::json::Value;
 
+using Clock = std::chrono::steady_clock;
+
 std::string json_quote(std::string_view value) {
     return engine::io::json::stringify_string(value);
 }
@@ -108,7 +112,49 @@ std::string base64_encode(const std::vector<uint8_t> & bytes) {
     return base64_encode(bytes.data(), bytes.size());
 }
 
-std::string task_result_json(const engine::runtime::TaskResult & result) {
+double elapsed_ms(Clock::time_point started) {
+    return std::chrono::duration<double, std::milli>(Clock::now() - started).count();
+}
+
+double audio_duration_ms(const engine::runtime::AudioBuffer & audio) {
+    if (audio.sample_rate <= 0 || audio.channels <= 0) {
+        return 0.0;
+    }
+    return 1000.0 * static_cast<double>(audio.samples.size()) /
+        static_cast<double>(audio.sample_rate * audio.channels);
+}
+
+double audio_rtf(double wall_ms, double duration_ms) {
+    return duration_ms > 0.0 ? wall_ms / duration_ms : 0.0;
+}
+
+std::string timing_json(double wall_ms) {
+    std::ostringstream out;
+    out << "{\"wall_ms\":" << wall_ms << "}";
+    return out.str();
+}
+
+std::string timing_json(double wall_ms, const engine::runtime::AudioBuffer & audio) {
+    const double duration_ms = audio_duration_ms(audio);
+    std::ostringstream out;
+    out << "{\"wall_ms\":" << wall_ms
+        << ",\"audio_duration_ms\":" << duration_ms
+        << ",\"rtf\":" << audio_rtf(wall_ms, duration_ms) << "}";
+    return out.str();
+}
+
+std::unordered_map<std::string, std::string> timing_headers(
+    double wall_ms,
+    const engine::runtime::AudioBuffer & audio) {
+    const double duration_ms = audio_duration_ms(audio);
+    return {
+        {"X-AudioCPP-Wall-Ms", std::to_string(wall_ms)},
+        {"X-AudioCPP-Audio-Duration-Ms", std::to_string(duration_ms)},
+        {"X-AudioCPP-RTF", std::to_string(audio_rtf(wall_ms, duration_ms))},
+    };
+}
+
+std::string task_result_json(const engine::runtime::TaskResult & result, double wall_ms) {
     std::ostringstream out;
     out << "{";
     bool first = true;
@@ -197,6 +243,14 @@ std::string task_result_json(const engine::runtime::TaskResult & result) {
         }
         out << "]";
     }
+    field("timing");
+    if (result.audio_output.has_value()) {
+        out << timing_json(wall_ms, *result.audio_output);
+    } else if (result.named_audio_outputs.size() == 1) {
+        out << timing_json(wall_ms, result.named_audio_outputs.front().audio);
+    } else {
+        out << timing_json(wall_ms);
+    }
     out << "}";
     return out.str();
 }
@@ -364,37 +418,55 @@ ServerState::LoadedModel & ServerState::require_model(const Value & body) {
     return *models_.at(it->second);
 }
 
-engine::runtime::TaskResult ServerState::run_model(
+struct ServerState::TimedTaskResult {
+    engine::runtime::TaskResult result;
+    double wall_ms = 0.0;
+};
+
+ServerState::TimedTaskResult ServerState::run_model(
     LoadedModel & model,
     const engine::runtime::TaskRequest & request) {
     std::lock_guard<std::mutex> lock(model.mutex);
     ensure_model_loaded_locked(model);
+    const auto started = Clock::now();
     model.session->prepare(engine::runtime::build_preparation_request(request));
-    return model.offline->run(request);
+    auto result = model.offline->run(request);
+    return TimedTaskResult{std::move(result), elapsed_ms(started)};
 }
 
 HttpResponse ServerState::handle_speech(const std::string & body_text) {
     const auto body = engine::io::json::parse(body_text);
     auto & model = require_model(body);
     const auto request = build_openai_speech_request(body, request_base_);
-    const auto result = run_model(model, request);
-    const auto wav = encode_pcm16_wav(select_audio_output(result));
+    const auto timed_result = run_model(model, request);
+    const auto & audio = select_audio_output(timed_result.result);
+    const auto wav = encode_pcm16_wav(audio);
     const auto response_format = engine::io::json::optional_string(body, "response_format", "wav");
     if (response_format == "json" || response_format == "b64_json") {
-        return json_response("{\"audio\":" + json_quote(base64_encode(wav)) + ",\"format\":\"wav\"}");
-    }
-    return HttpResponse{200, "audio/wav", std::string(reinterpret_cast<const char *>(wav.data()), wav.size()), {}};
+        return json_response(
+            "{\"audio\":" + json_quote(base64_encode(wav)) +
+            ",\"format\":\"wav\",\"timing\":" + timing_json(timed_result.wall_ms, audio) + "}");
+    }
+    return HttpResponse{
+        200,
+        "audio/wav",
+        std::string(reinterpret_cast<const char *>(wav.data()), wav.size()),
+        timing_headers(timed_result.wall_ms, audio),
+    };
 }
 
 HttpResponse ServerState::handle_transcription(const std::string & body_text) {
     const auto body = engine::io::json::parse(body_text);
     auto & model = require_model(body);
     const auto request = build_openai_transcription_request(body, request_base_);
-    const auto result = run_model(model, request);
+    const auto timed_result = run_model(model, request);
+    const auto & result = timed_result.result;
     if (!result.text_output.has_value()) {
         throw std::runtime_error("model result did not contain transcript text");
     }
-    return json_response("{\"text\":" + json_quote(result.text_output->text) + "}");
+    return json_response(
+        "{\"text\":" + json_quote(result.text_output->text) +
+        ",\"timing\":" + timing_json(timed_result.wall_ms) + "}");
 }
 
 HttpResponse ServerState::handle_generic_run(const std::string & body_text) {
@@ -404,7 +476,8 @@ HttpResponse ServerState::handle_generic_run(const std::string & body_text) {
     const auto request = minitts::cli::build_request_from_json(
         request_json != nullptr ? *request_json : body,
         request_base_);
-    return json_response(task_result_json(run_model(model, request)));
+    const auto timed_result = run_model(model, request);
+    return json_response(task_result_json(timed_result.result, timed_result.wall_ms));
 }
 
 std::string ServerState::models_json() const {

app/server/runtime.h

@@ -34,7 +34,8 @@ class ServerState final : public IHttpHandler {
     void load_models();
     void ensure_model_loaded_locked(LoadedModel & model);
     LoadedModel & require_model(const engine::io::json::Value & body);
-    engine::runtime::TaskResult run_model(LoadedModel & model, const engine::runtime::TaskRequest & request);
+    struct TimedTaskResult;
+    TimedTaskResult run_model(LoadedModel & model, const engine::runtime::TaskRequest & request);
     HttpResponse handle_speech(const std::string & body_text);
     HttpResponse handle_transcription(const std::string & body_text);
     HttpResponse handle_generic_run(const std::string & body_text);

docs/tts.md

@@ -42,7 +42,7 @@ audiocpp_cli --task clon --family chatterbox --model models/chatterbox --backend
 |---|---|---:|---|
 | `--voice-ref` | WAV path | required | Reference speaker audio. |
 | `--language` | language code | `en` | Text language. |
-| `--text-chunk-size` | integer chars | `256` | Long-form chunk size. |
+| `--text-chunk-size` | integer chars | `128` | Long-form chunk size. |
 | `--guidance-scale` | float | `0.5` | CFG strength. |
 | `--temperature` | float | `0.8` | T3 sampling temperature. |
 | `--top-p` | float | `0.8` | T3 nucleus sampling limit. |

include/engine/framework/modules/hift_vocoder.h

@@ -144,6 +144,7 @@ class HiftVocoderComponent {
         uint64_t prior_noise_values = 0,
         const std::vector<float> * source_random_values = nullptr) const;
     std::vector<float> predict_f0(const std::vector<float> & mel, int64_t frames) const;
+    void release_runtime_cache() const;
 
 private:
     struct State;

include/engine/framework/runtime/cache_slots.h

@@ -0,0 +1,112 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <utility>
+#include <vector>
+
+namespace engine::runtime {
+
+template <typename Key, typename Value, typename Equal = std::equal_to<Key>>
+class CacheSlots {
+public:
+    explicit CacheSlots(std::size_t capacity = 1)
+        : capacity_(capacity) {}
+
+    std::size_t capacity() const noexcept {
+        return capacity_;
+    }
+
+    std::size_t size() const noexcept {
+        return entries_.size();
+    }
+
+    void set_capacity(std::size_t capacity) {
+        capacity_ = capacity;
+        evict_to_capacity();
+    }
+
+    Value * find(const Key & key) {
+        for (auto & entry : entries_) {
+            if (equal_(entry.key, key)) {
+                entry.last_used = next_tick();
+                return &entry.value;
+            }
+        }
+        return nullptr;
+    }
+
+    const Value * find(const Key & key) const {
+        for (const auto & entry : entries_) {
+            if (equal_(entry.key, key)) {
+                return &entry.value;
+            }
+        }
+        return nullptr;
+    }
+
+    void put(Key key, Value value) {
+        if (capacity_ == 0) {
+            entries_.clear();
+            return;
+        }
+        for (auto & entry : entries_) {
+            if (equal_(entry.key, key)) {
+                entry.key = std::move(key);
+                entry.value = std::move(value);
+                entry.last_used = next_tick();
+                return;
+            }
+        }
+        if (entries_.size() >= capacity_) {
+            erase_lru();
+        }
+        entries_.push_back(Entry{
+            std::move(key),
+            std::move(value),
+            next_tick(),
+        });
+    }
+
+    void clear() {
+        entries_.clear();
+    }
+
+private:
+    struct Entry {
+        Key key;
+        Value value;
+        std::uint64_t last_used = 0;
+    };
+
+    std::uint64_t next_tick() noexcept {
+        return ++tick_;
+    }
+
+    void evict_to_capacity() {
+        while (entries_.size() > capacity_) {
+            erase_lru();
+        }
+    }
+
+    void erase_lru() {
+        if (entries_.empty()) {
+            return;
+        }
+        std::size_t oldest = 0;
+        for (std::size_t i = 1; i < entries_.size(); ++i) {
+            if (entries_[i].last_used < entries_[oldest].last_used) {
+                oldest = i;
+            }
+        }
+        entries_.erase(entries_.begin() + static_cast<std::ptrdiff_t>(oldest));
+    }
+
+    std::size_t capacity_ = 1;
+    std::uint64_t tick_ = 0;
+    std::vector<Entry> entries_;
+    Equal equal_;
+};
+
+}  // namespace engine::runtime

include/engine/models/chatterbox/components.h

@@ -199,6 +199,7 @@ class HiFTVocoderComponent {
         uint64_t seed,
         uint64_t prior_noise_values,
         const std::vector<float> & cache_source) const;
+    void release_runtime_cache() const;
 
 private:
     struct State;

include/engine/models/chatterbox/s3gen_inference.h

@@ -22,6 +22,8 @@ class S3GenSessionCache {
     S3GenSessionCache(const S3GenSessionCache &) = delete;
     S3GenSessionCache & operator=(const S3GenSessionCache &) = delete;
 
+    void release_runtime_graphs();
+
 private:
     struct State;
     std::unique_ptr<State> state_;

include/engine/models/chatterbox/session.h

@@ -2,16 +2,31 @@
 
 #include "engine/framework/runtime/session_base.h"
 #include "engine/framework/assets/tensor_source.h"
+#include "engine/framework/runtime/cache_slots.h"
 #include "engine/models/chatterbox/assets.h"
 #include "engine/models/chatterbox/conditionals.h"
 #include "engine/models/chatterbox/tts.h"
 
+#include <cstddef>
 #include <filesystem>
 #include <memory>
 #include <optional>
+#include <string>
 
 namespace engine::models::chatterbox {
 
+struct ChatterboxConditionalsCacheKey {
+    runtime::AudioBuffer reference_audio;
+    float exaggeration = 0.0f;
+    std::string language;
+};
+
+struct ChatterboxConditionalsCacheKeyEqual {
+    bool operator()(
+        const ChatterboxConditionalsCacheKey & lhs,
+        const ChatterboxConditionalsCacheKey & rhs) const;
+};
+
 class ChatterboxSession final
     : public runtime::RuntimeSessionBase
     , public runtime::IOfflineVoiceTaskSession {
@@ -35,8 +50,14 @@ class ChatterboxSession final
     std::shared_ptr<const ChatterboxAssetPaths> assets_;
     engine::assets::TensorStorageType t3_weight_storage_type_ = engine::assets::TensorStorageType::Native;
     engine::assets::TensorStorageType component_weight_storage_type_ = engine::assets::TensorStorageType::Native;
+    bool mem_saver_ = false;
     std::unique_ptr<ChatterboxTtsComponent> component_;
+    std::optional<std::string> component_language_;
     std::optional<ChatterboxVoiceCloneConfig> voice_clone_config_;
+    runtime::CacheSlots<
+        ChatterboxConditionalsCacheKey,
+        ChatterboxConditionalsOutputs,
+        ChatterboxConditionalsCacheKeyEqual> conditionals_cache_;
     std::optional<ChatterboxConditionalsOutputs> cached_conditionals_;
     double cached_prompt_prep_ms_ = 0.0;
 };