chatterbox: cache prepared voice-conversion targets

docs/tts.md

@@ -152,6 +152,7 @@ settings for subsequent requests in that session.
 | `chatterbox.t3_weight_type` | `native`, `f32`, `f16`, `bf16`, `q8_0` | shared setting | T3 weight storage; inherits weight_type. |
 | `chatterbox.multilingual_t3` | `v2`, `v3` | `v2` | Multilingual T3 checkpoint selection. |
 | `chatterbox.conditionals_cache_slots` | integer >= 0 | `1` | Cached speaker conditionals; zero disables reuse. |
+| `chatterbox.vc_reference_cache_slots` | integer >= 0 | `1` | Cached voice-conversion target references (`--task vc`), so converting into the same voice again skips preparing it; zero disables reuse. |
 | `chatterbox.mem_saver` | bool | `false` | Release staged runtime components. |
 | `chatterbox.encoder_condition_samples` | integer | `96000` | Tokenizer reference sample limit at 16 kHz; the voice encoder uses the full reference. |
 | `chatterbox.decoder_condition_samples` | integer | `240000` | S3Gen reference sample limit at 24 kHz. |

include/engine/models/chatterbox/session.h

@@ -27,6 +27,10 @@ struct ChatterboxConditionalsCacheKey {
     std::string language;
 };
 
+struct ChatterboxAudioBufferEqual {
+    bool operator()(const runtime::AudioBuffer & lhs, const runtime::AudioBuffer & rhs) const;
+};
+
 struct ChatterboxConditionalsCacheKeyEqual {
     bool operator()(
         const ChatterboxConditionalsCacheKey & lhs,
@@ -68,6 +72,10 @@ class ChatterboxSession final
         ChatterboxConditionalsOutputs,
         ChatterboxConditionalsCacheKeyEqual> conditionals_cache_;
     std::optional<ChatterboxConditionalsOutputs> cached_conditionals_;
+    runtime::CacheSlots<
+        runtime::AudioBuffer,
+        EmbedReferenceOutputs,
+        ChatterboxAudioBufferEqual> vc_reference_cache_;
     double cached_prompt_prep_ms_ = 0.0;
 };
 

include/engine/models/chatterbox/vc.h

@@ -46,6 +46,20 @@ class ChatterboxVCComponent {
         const runtime::AudioBuffer & target_voice,
         const ChatterboxVoiceConversionConfig & config = {}) const;
 
+    // Rejects invalid settings or source audio, before any target is prepared.
+    void validate(
+        const runtime::AudioBuffer & source_audio,
+        const ChatterboxVoiceConversionConfig & config) const;
+
+    // The target speaker's prepared reference, reusable for every
+    // conversion into that voice.
+    EmbedReferenceOutputs prepare_target(const runtime::AudioBuffer & target_voice) const;
+
+    ChatterboxVoiceConversionOutputs convert(
+        const runtime::AudioBuffer & source_audio,
+        const EmbedReferenceOutputs & target_ref,
+        const ChatterboxVoiceConversionConfig & config = {}) const;
+
 private:
     engine::models::chatterbox::S3TokenizerComponent tokenizer_;
     engine::models::chatterbox::CAMPPlusEncoderComponent speaker_encoder_;

model_specs/chatterbox.json

@@ -176,6 +176,14 @@
         "min": 0,
         "default": 1
       },
+      {
+        "name": "vc_reference_cache_slots",
+        "type": "int",
+        "description": "Cached voice-conversion target references; zero disables reuse.",
+        "required": false,
+        "min": 0,
+        "default": 1
+      },
       {
         "name": "mem_saver",
         "type": "bool",

src/models/chatterbox/loader.cpp

@@ -55,6 +55,7 @@ runtime::ModelCliInterface cli(const ChatterboxAssets &) {
         {"chatterbox.t3_weight_type", "native|f32|f16|bf16|q8_0", "T3 weight storage; inherits weight_type.", false},
         {"chatterbox.multilingual_t3", "v2|v3", "Multilingual T3 checkpoint for non-English TTS; default v2.", false, "v2"},
         {"chatterbox.conditionals_cache_slots", "int", "Prepared voice-condition cache slots; default 1, set 0 to disable.", false, "1", "0"},
+        {"chatterbox.vc_reference_cache_slots", "int", "Prepared voice-conversion target cache slots; default 1, set 0 to disable.", false, "1", "0"},
         {"chatterbox.mem_saver", "bool", "Free non-conditional runtime graphs after each request chunk; default false.", false, "false"},
         {"chatterbox.encoder_condition_samples", "int", "Maximum reference samples for the T3 tokenizer; the voice encoder still consumes the full 16 kHz reference. Default: 96000.", false, "96000"},
         {"chatterbox.decoder_condition_samples", "int", "S3Gen reference sample limit at 24 kHz.", false, "240000"},

src/models/chatterbox/session.cpp

@@ -263,6 +263,7 @@ void validate_chatterbox_options(const runtime::SessionOptions & options) {
             key != "chatterbox.t3_weight_type" &&
             key != "chatterbox.multilingual_t3" &&
             key != "chatterbox.conditionals_cache_slots" &&
+            key != "chatterbox.vc_reference_cache_slots" &&
             key != "chatterbox.mem_saver" &&
             key != "chatterbox.encoder_condition_samples" &&
             key != "chatterbox.decoder_condition_samples" &&
@@ -318,6 +319,17 @@ std::size_t resolve_conditionals_cache_slots(const runtime::SessionOptions & opt
     return static_cast<std::size_t>(slots);
 }
 
+std::size_t resolve_vc_reference_cache_slots(const runtime::SessionOptions & options) {
+    const auto slots = runtime::parse_i64_option(
+        options.options,
+        {"chatterbox.vc_reference_cache_slots", "vc_reference_cache_slots"})
+        .value_or(1);
+    if (slots < 0) {
+        throw std::runtime_error("chatterbox.vc_reference_cache_slots must be non-negative");
+    }
+    return static_cast<std::size_t>(slots);
+}
+
 bool resolve_mem_saver(const runtime::SessionOptions & options) {
     if (const auto value = runtime::find_option(options.options, {"chatterbox.mem_saver", "mem_saver"})) {
         return runtime::parse_bool_option(*value, "chatterbox.mem_saver");
@@ -327,6 +339,12 @@ bool resolve_mem_saver(const runtime::SessionOptions & options) {
 
 }  // namespace
 
+bool ChatterboxAudioBufferEqual::operator()(
+    const runtime::AudioBuffer & lhs,
+    const runtime::AudioBuffer & rhs) const {
+    return same_audio_buffer(lhs, rhs);
+}
+
 bool ChatterboxConditionalsCacheKeyEqual::operator()(
     const ChatterboxConditionalsCacheKey & lhs,
     const ChatterboxConditionalsCacheKey & rhs) const {
@@ -346,7 +364,8 @@ ChatterboxSession::ChatterboxSession(
       component_weight_storage_type_(resolve_component_weight_storage_type(this->options())),
       multilingual_t3_version_(resolve_multilingual_t3_version(this->options())),
       mem_saver_(resolve_mem_saver(this->options())),
-      conditionals_cache_(resolve_conditionals_cache_slots(this->options())) {
+      conditionals_cache_(resolve_conditionals_cache_slots(this->options())),
+      vc_reference_cache_(resolve_vc_reference_cache_slots(this->options())) {
     if (task_.task != runtime::VoiceTaskKind::Tts &&
         task_.task != runtime::VoiceTaskKind::VoiceCloning &&
         task_.task != runtime::VoiceTaskKind::VoiceConversion) {
@@ -476,7 +495,21 @@ runtime::TaskResult ChatterboxSession::run_voice_conversion(const runtime::TaskR
     }
 
     const auto config = make_voice_conversion_config(request.options);
-    auto outputs = vc_component_->convert(source_audio, target_voice, config);
+    vc_component_->validate(source_audio, config);
+    ChatterboxVoiceConversionOutputs outputs;
+    if (const auto * target_ref = vc_reference_cache_.find(target_voice)) {
+        engine::debug::trace_log_scalar("chatterbox.vc.reference_cache_hit", 1);
+        outputs = vc_component_->convert(source_audio, *target_ref, config);
+    } else {
+        engine::debug::trace_log_scalar("chatterbox.vc.reference_cache_hit", 0);
+        const auto target_ref_start = std::chrono::steady_clock::now();
+        auto prepared = vc_component_->prepare_target(target_voice);
+        const double target_ref_ms = engine::debug::elapsed_ms(target_ref_start);
+        engine::debug::timing_log_scalar("chatterbox.vc.target_ref_ms", target_ref_ms);
+        outputs = vc_component_->convert(source_audio, prepared, config);
+        outputs.target_ref_ms = target_ref_ms;
+        vc_reference_cache_.put(std::move(target_voice), std::move(prepared));
+    }
     runtime::TaskResult result;
     result.audio_output = runtime::AudioBuffer{
         24000,

src/models/chatterbox/vc.cpp

@@ -40,6 +40,15 @@ runtime::AudioBuffer trim_mono_audio(
     return audio;
 }
 
+void validate_config(const ChatterboxVoiceConversionConfig & config) {
+    if (config.num_inference_steps <= 0) {
+        throw std::runtime_error("Chatterbox VC num_inference_steps must be positive");
+    }
+    if (!(config.s3gen_cfg_rate >= 0.0F)) {
+        throw std::runtime_error("Chatterbox VC s3gen_cfg_rate must be non-negative");
+    }
+}
+
 }  // namespace
 
 runtime::AudioBuffer load_chatterbox_vc_audio_mono(
@@ -84,28 +93,40 @@ ChatterboxVoiceConversionOutputs ChatterboxVCComponent::convert(
     const runtime::AudioBuffer & source_audio,
     const runtime::AudioBuffer & target_voice,
     const ChatterboxVoiceConversionConfig & config) const {
-    if (config.num_inference_steps <= 0) {
-        throw std::runtime_error("Chatterbox VC num_inference_steps must be positive");
-    }
-    if (!(config.s3gen_cfg_rate >= 0.0F)) {
-        throw std::runtime_error("Chatterbox VC s3gen_cfg_rate must be non-negative");
-    }
+    validate(source_audio, config);
+    const auto target_ref_start = std::chrono::steady_clock::now();
+    const auto target_ref = prepare_target(target_voice);
+    const double target_ref_ms = engine::debug::elapsed_ms(target_ref_start);
+    auto outputs = convert(source_audio, target_ref, config);
+    outputs.target_ref_ms = target_ref_ms;
+    engine::debug::timing_log_scalar("chatterbox.vc.target_ref_ms", target_ref_ms);
+    return outputs;
+}
+
+void ChatterboxVCComponent::validate(
+    const runtime::AudioBuffer & source_audio,
+    const ChatterboxVoiceConversionConfig & config) const {
+    validate_config(config);
     validate_audio(source_audio, "source");
-    validate_audio(target_voice, "target");
+}
 
-    ChatterboxVoiceConversionOutputs outputs;
-    const auto source_16k = normalize_chatterbox_vc_audio_mono(source_audio, 16000);
+EmbedReferenceOutputs ChatterboxVCComponent::prepare_target(const runtime::AudioBuffer & target_voice) const {
+    validate_audio(target_voice, "target");
     const auto target_24k = trim_mono_audio(
         normalize_chatterbox_vc_audio_mono(target_voice, 24000),
         prompt_prep_config_.decoder_condition_samples);
     const auto target_16k = normalize_chatterbox_vc_audio_mono(target_24k, 16000);
+    return tokenizer_.embed_reference_from_wavs(speaker_encoder_, target_24k, target_16k);
+}
 
-    const auto target_ref_start = std::chrono::steady_clock::now();
-    auto target_ref = tokenizer_.embed_reference_from_wavs(
-        speaker_encoder_,
-        target_24k,
-        target_16k);
-    outputs.target_ref_ms = engine::debug::elapsed_ms(target_ref_start);
+ChatterboxVoiceConversionOutputs ChatterboxVCComponent::convert(
+    const runtime::AudioBuffer & source_audio,
+    const EmbedReferenceOutputs & target_ref,
+    const ChatterboxVoiceConversionConfig & config) const {
+    validate(source_audio, config);
+
+    ChatterboxVoiceConversionOutputs outputs;
+    const auto source_16k = normalize_chatterbox_vc_audio_mono(source_audio, 16000);
 
     const auto source_tokenizer_start = std::chrono::steady_clock::now();
     auto source_tokens = tokenizer_.tokenize(source_16k, std::nullopt);
@@ -139,7 +160,6 @@ ChatterboxVoiceConversionOutputs ChatterboxVCComponent::convert(
         s3gen_->release_vocoder_graphs();
     }
 
-    engine::debug::timing_log_scalar("chatterbox.vc.target_ref_ms", outputs.target_ref_ms);
     engine::debug::timing_log_scalar("chatterbox.vc.source_tokenizer_ms", outputs.source_tokenizer_ms);
     engine::debug::timing_log_scalar("chatterbox.vc.s3gen.total_ms", outputs.s3gen_ms);
     engine::debug::timing_log_scalar("chatterbox.vc.s3gen.token2mel.total_ms", outputs.s3gen_token2mel_ms);