Echo-TTS: Implement optional Joint-CFG

docs/community_models/echo_tts_performance.md

@@ -118,6 +118,45 @@ Verified against a mock denoiser at 0.011% max deviation for interval 2 and
 0.020% for interval 3 -- but that mock had deliberately smooth uncond offsets,
 so treat those as a lower bound, not a measurement on real audio.
 
+### 5. `echo_tts.mem_saver` (opt-in, default false)
+Mem-saver frees VRAM by free-ing the audio encoder allocations after reference
+speaker encoding. For cached reference speakers, this has no performance hit,
+for an uncached reference speaker this has a small performance hit. Peak VRAM
+often occurs during speaker encoding, but depends on the length of prompts you
+send (long prompts make the audio decode buffers larger than the encode buffers).
+Actual savings will depend on model quantization and backend architecture. For
+Q8_0 model weights, the following were measured on a GTX 1080 Ti and a RTX 4070:
+- GTX 1080 Ti (false/true): 8122 MB / 7416 MB (8.7% savings)
+- RTX 4070  (false/true): 6084 MB / 5356 MB (12.0% savings)
+
+### 6. `echo_tts.cfg_mode` (opt-in, default independent)
+Echo-TTS allows two classifier free guidance modes: the reference Independent
+mode with a guidance scale for the reference speaker and a guidance scale for the
+text, and a Joint mode with a single combined guidance scale. Joint-CFG allows
+for a slightly smaller VRAM allocation, and fewer forwards passes. While actual
+results will depend on the backend architecture, the following numbers are
+provided as examples measured on a GTX 1080 Ti and an RTX 4080:
+- GTX 1080 Ti (independent/joint): 16.9s / 11.9s (0.7x, or ~30% faster)
+- RTX 4070 (independent/joint): 3.5s / 2.7s (0.77x, or ~23% faster)
+
+The setting also has an effect on audio quality. Measured by @5uck1ess on an RTX
+5090, 4 English prompts x 3 seeds run through Whisper-v3 WER and WavLM SIM:
+
+| Mode & Scale | WER | UTMOS | SIM |
+|---|---|---|---|
+| independent, 3.0/8.0 | 0.069 | 4.20 | 0.831 |
+| joint, 4.0 | 0.081 | 4.32 | 0.804 |
+| joint, 6.0 | 0.060 | 4.31 | 0.821 |
+| joint, 8.0 | 0.076 | 4.29 | 0.829 |
+| joint, 10.0 | 0.069 | 4.26 | 0.830 |
+
+Joint-CFG at scale 8.0 is at parity, below ~6.0 scale similarity drops. By ear:
+near parity, independent is slightly better on a few, clone is almost identical.
+@lojack5 also notes: Joint-CFG the most noticable (by ear) difference is at the
+first ~0.2 seconds of the generated audio, where the speaker doesn't as closely
+match the reference as Independent-CFG does. After the beginning of the
+generation, the speaker quickly matches the reference audio.
+
 ---
 
 ## Suggested config

include/engine/community_models/echo_tts/config.h

@@ -161,11 +161,19 @@ struct EchoPcaState {
     float latent_scale = 1.0F;
 };
 
+// Which CFG method to use when sampling (may be expanded in the future)
+enum class EchoCfgMode {
+    Independent,    // Reference independent 2-CFG (speaker, text) method
+    Joint,          // Single CFG (speaker+text) method
+};
+
 // Request-level sampler configuration, parsed from spec options.
 struct EchoSamplerOptions {
     int num_steps = 40;
-    float cfg_scale_text = 3.0F;
-    float cfg_scale_speaker = 8.0F;
+    EchoCfgMode cfg_mode = EchoCfgMode::Independent;
+    float cfg_scale_text = 3.0F;        // Independent-CFG only
+    float cfg_scale_speaker = 8.0F;     // Independent-CFG only
+    float cfg_scale = 8.0F;             // Joint-CFG only
     float cfg_min_t = 0.5F;
     float cfg_max_t = 1.0F;
     // Evaluate the two unconditional lanes only every Nth step inside the CFG

include/engine/community_models/echo_tts/dit.h

@@ -47,7 +47,7 @@ class EchoDiTRuntime {
 
     // Runs the text and speaker encoders and populates the cached key/value
     // projections. Must be called before sample().
-    void prepare_conditioning(const EchoConditioning & conditioning);
+    void prepare_conditioning(const EchoConditioning & conditioning, int max_kv_lanes);
 
     // Runs the dual-CFG Euler sampler and returns the final latent, shaped
     // (sequence_length, latent_size) row-major.

include/engine/community_models/echo_tts/sampler.h

@@ -31,12 +31,21 @@ bool cfg_active(float t, float cfg_min_t, float cfg_max_t);
 //     v = v_cond
 //         + w_text    * (v_cond - v_uncond_text)
 //         + w_speaker * (v_cond - v_uncond_speaker)
-std::vector<float> combine_cfg_lanes(
+std::vector<float> combine_cfg_lanes_independent(
     const std::vector<float> & lanes,
     int64_t lane_elements,
     float cfg_scale_text,
     float cfg_scale_speaker);
 
+// Combines the two CFG lanes into a single velocity, following the
+// joint-guidance form:
+//     v = v_cond
+//         + w_joint * (v_cond - v_undcond)
+std::vector<float> combine_cfg_lanes_joint(
+    const std::vector<float> & lanes,
+    int64_t lane_elements,
+    float cfg_scale);
+
 // Runs the sampler loop. `denoise` supplies the model evaluation and
 // `initial_noise` the starting latent, both injected so this can be tested
 // without a backend. `on_kv_rescale`, when set, is invoked at the timestep where

include/engine/community_models/echo_tts/session.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "config.h"
 #include "engine/community_models/echo_tts/config.h"
 #include "engine/framework/assets/resource_bundle.h"
 #include "engine/framework/model_spec/metadata.h"
@@ -86,6 +87,7 @@ class EchoTtsSession final
     std::shared_ptr<const engine::model_spec::ModelContract> contract_;
     std::unique_ptr<EchoDiTRuntime> dit_;
     std::unique_ptr<engine::codecs::FishDacCodecRuntime> codec_;
+    EchoCfgMode cfg_mode_ = EchoCfgMode::Independent;
     int64_t reference_max_samples_ = 0;
     bool mem_saver_ = false;
     std::vector<float> speaker_latent_;

model_specs/echo_tts.json

@@ -38,15 +38,23 @@
       {
         "name": "text_guidance_scale",
         "type": "float",
-        "description": "Classifier-free guidance scale on the text condition. Advertised range is recommended, not enforced by the parser.",
+        "description": "Classifier-free guidance scale on the text condition (independent-cfg only). Advertised range is recommended, not enforced by the parser.",
         "required": false,
         "min": 0,
         "default": 3
       },
       {
         "name": "speaker_guidance_scale",
         "type": "float",
-        "description": "Classifier-free guidance scale on the speaker condition. Advertised range is recommended, not enforced by the parser.",
+        "description": "Classifier-free guidance scale on the speaker condition (independent-cfg only). Advertised range is recommended, not enforced by the parser.",
+        "required": false,
+        "min": 0,
+        "default": 8
+      },
+      {
+        "name": "guidance_scale",
+        "type": "float",
+        "description": "Classifier-free guidance scale on the text and speaker condition (join-cfg only). Advertised range is recommended, not enforced by the parser.",
         "required": false,
         "min": 0,
         "default": 8
@@ -126,6 +134,17 @@
         "required": false,
         "description": "Release Fish S1-DAC decode graph space before encoding reference speaker audio and encode graph space after. Avoids having both allocations resident simulaneously at the cost of rebuilding the graphs when a speaker cache-miss occurs.",
         "default": false
+      },
+      {
+        "name": "cfg_mode",
+        "type": "enum",
+        "description": "Type of classifier free guidance used.",
+        "values": [
+          "independent",
+          "joint"
+        ],
+        "required": false,
+        "default": "independent"
       }
     ],
     "load": []

src/community_models/echo_tts/dit.cpp

@@ -214,7 +214,7 @@ class EchoDiTRuntime::Impl {
     const EchoTtsConfig & config() const noexcept { return config_; }
     bool conditioning_ready() const noexcept { return conditioning_ready_; }
 
-    void prepare_conditioning(const EchoConditioning & conditioning) {
+    void prepare_conditioning(const EchoConditioning & conditioning, int max_kv_lanes) {
         validate_conditioning(conditioning);
 
         // Any change in conditioning length invalidates every cached graph,
@@ -235,7 +235,7 @@ class EchoDiTRuntime::Impl {
                 conditioning.speaker_mask[static_cast<size_t>(i * config_.speaker_patch_size)];
         }
 
-        allocate_kv_cache();
+        allocate_kv_cache(max_kv_lanes);
         build_conditioning_graph();
 
         core::write_tensor_i32(text_ids_, conditioning.text_input_ids);
@@ -281,15 +281,15 @@ class EchoDiTRuntime::Impl {
         if (!conditioning_ready_) {
             throw std::runtime_error("Echo-TTS denoise() called before prepare_conditioning()");
         }
-        if (lanes != 1 && lanes != 3) {
-            throw std::runtime_error("Echo-TTS denoiser supports 1 or 3 CFG lanes");
+        if (lanes != 1 && lanes != 2 && lanes != 3) {
+            throw std::runtime_error("Echo-TTS denoiser supports 1, 2, or 3 CFG lanes");
         }
         const int64_t elements = sequence_length_ * config_.latent_size;
         if (static_cast<int64_t>(x.size()) != elements) {
             throw std::runtime_error("Echo-TTS denoiser received a mis-shaped latent");
         }
 
-        auto & graph = lanes == 1 ? single_ : triple_;
+        auto & graph = lanes == 1 ? single_ : lanes == 2 ? joint_ : independent_;
         if (graph.graph == nullptr) {
             build_denoiser_graph(graph, lanes);
         }
@@ -453,15 +453,18 @@ class EchoDiTRuntime::Impl {
     }
 
     // Denoiser mask, laid out per lane as [self | text | speaker]. Lane 0 is
-    // fully conditional, lane 1 drops text, lane 2 drops speaker, reproducing
-    // upstream's concatenated cond/uncond masks.
+    // fully conditional,
+    // joint-cfg: lane 1 drops text and speaker
+    // independent-cfg: lane 1 drops text, lane 2 drops speaker, reproducing
+    //     upstream's concatenated cond/uncond masks.
     std::vector<float> make_denoiser_mask(int lanes) const {
         const int64_t keys = sequence_length_ + text_length_ + speaker_tokens_;
         std::vector<float> mask(
             static_cast<size_t>(static_cast<int64_t>(lanes) * sequence_length_ * keys), 0.0F);
         for (int lane = 0; lane < lanes; ++lane) {
+            // Note: logic needs to be correct for both lanes == 1 and lanes == 2/3
             const bool text_on = lane != 1;
-            const bool speaker_on = lane != 2;
+            const bool speaker_on = lane == 0 || lane != (lanes - 1);
             for (int64_t q = 0; q < sequence_length_; ++q) {
                 float * row = mask.data() +
                               (static_cast<int64_t>(lane) * sequence_length_ + q) * keys;
@@ -490,7 +493,7 @@ class EchoDiTRuntime::Impl {
     // The cache lives in its own context and backend buffer so that both the
     // conditioning graph (which writes it) and the denoiser graphs (which read
     // it) can reference the same tensors as leaves.
-    void allocate_kv_cache() {
+    void allocate_kv_cache(int max_kv_lanes) {
         const int64_t heads = config_.num_heads;
         const int64_t head_dim = config_.head_dim();
         const size_t tensor_count = static_cast<size_t>(config_.num_layers) * 4;
@@ -501,12 +504,13 @@ class EchoDiTRuntime::Impl {
             throw std::runtime_error("Echo-TTS KV cache context initialization failed");
         }
         core::ModuleBuildContext ctx{kv_ctx_.get(), "echo_tts.kv_cache", backend_type_};
-        // The cache is allocated at the widest lane count any graph will ask
-        // for, so joint_attention's expand() finds dims[0] already equal to the
+        // The cache is allocated at the widest lane count any graph will ask for in the
+        // session-configured CFG mode,
+        // so joint_attention's expand() finds dims[0] already equal to the
         // batch and short-circuits to a passthrough. The single-lane graph
         // reads lane 0, which is a contiguous prefix because the lane axis is
         // outermost in ggml's layout.
-        kv_lanes_ = echo_kv_expand_disabled() ? 1 : kMaxCfgLanes;
+        kv_lanes_ = echo_kv_expand_disabled() ? 1 : max_kv_lanes;
         auto make = [&](int64_t tokens) {
             return core::make_tensor(
                 ctx, GGML_TYPE_F32,
@@ -651,12 +655,12 @@ class EchoDiTRuntime::Impl {
 
     // Reconciles the cache's lane count with the graph's.
     //
-    // With pre-expansion on, the cache is allocated at kMaxCfgLanes and the
+    // With pre-expansion on, the cache is allocated at 2 or 3 lanes and the
     // single-lane graph reads a prefix of it; the lane axis is outermost, so
     // that is a view with no copy.
     //
     // With pre-expansion off (AUDIOCPP_ECHO_TTS_NO_KV_EXPAND=1) the cache is
-    // batch-1 and the three-lane graph is *wider* than it. Narrowing is not
+    // batch-1 and the two- or three-lane graphs are *wider* than it. Narrowing is not
     // possible and not wanted: returning the cache unchanged leaves
     // joint_attention's expand() to broadcast it per step, which is exactly the
     // pre-patch behaviour the flag exists to restore.
@@ -779,7 +783,8 @@ class EchoDiTRuntime::Impl {
 
     void release_denoiser_graphs() {
         release_denoiser_graph(single_);
-        release_denoiser_graph(triple_);
+        release_denoiser_graph(joint_);
+        release_denoiser_graph(independent_);
     }
 
     void release_kv_cache() {
@@ -831,7 +836,8 @@ class EchoDiTRuntime::Impl {
     core::TensorValue speaker_positions_;
 
     DenoiserGraph single_;
-    DenoiserGraph triple_;
+    DenoiserGraph joint_;
+    DenoiserGraph independent_;
 };
 
 EchoDiTRuntime::EchoDiTRuntime(
@@ -846,8 +852,8 @@ EchoDiTRuntime::~EchoDiTRuntime() = default;
 
 const EchoTtsConfig & EchoDiTRuntime::config() const noexcept { return impl_->config(); }
 
-void EchoDiTRuntime::prepare_conditioning(const EchoConditioning & conditioning) {
-    impl_->prepare_conditioning(conditioning);
+void EchoDiTRuntime::prepare_conditioning(const EchoConditioning & conditioning, int max_kv_lanes) {
+    impl_->prepare_conditioning(conditioning, max_kv_lanes);
 }
 
 std::vector<float> EchoDiTRuntime::denoise_once(const std::vector<float> & x, float t, int lanes) {

src/community_models/echo_tts/dit_blocks.inc

@@ -34,10 +34,6 @@ constexpr size_t kWeightContextBytes = 6144ull * 1024ull * 1024ull;
 constexpr size_t kGraphArenaBytes = 512ull * 1024ull * 1024ull;
 constexpr float kRopeTheta = 10000.0F;
 
-// cond, text-uncond, speaker-uncond. The KV cache is sized to this so the
-// three-lane graph never has to broadcast it at sampling time.
-constexpr int64_t kMaxCfgLanes = 3;
-
 // Additive mask value for disallowed keys. -INFINITY would be exact but
 // produces NaN when an entire row is masked, which happens for the
 // speaker-uncond CFG lane; a large finite penalty is the standard ggml

src/community_models/echo_tts/sampler.cpp

@@ -30,7 +30,7 @@ bool cfg_active(float t, float cfg_min_t, float cfg_max_t) {
     return t >= cfg_min_t && t <= cfg_max_t;
 }
 
-std::vector<float> combine_cfg_lanes(
+std::vector<float> combine_cfg_lanes_independent(
     const std::vector<float> & lanes,
     int64_t lane_elements,
     float cfg_scale_text,
@@ -53,6 +53,26 @@ std::vector<float> combine_cfg_lanes(
     return out;
 }
 
+std::vector<float> combine_cfg_lanes_joint(
+    const std::vector<float> & lanes,
+    int64_t lane_elements,
+    float cfg_scale) {
+    if (lane_elements <= 0 ||
+        static_cast<int64_t>(lanes.size()) != lane_elements * 2) {
+        throw std::runtime_error("Echo-TTS Joint-CFG combine expects exactly 2 lanes");
+    }
+    const float * cond = lanes.data();
+    const float * uncond = lanes.data() + lane_elements;
+
+    std::vector<float> out(static_cast<size_t>(lane_elements));
+    for (int64_t i = 0; i < lane_elements; ++i) {
+        const float c = cond[i];
+        out[static_cast<size_t>(i)] =
+            c + cfg_scale * (c - uncond[i]);
+    }
+    return out;
+}
+
 std::vector<float> run_euler_sampler(
     const EchoSamplerOptions & options,
     int64_t sequence_length,
@@ -84,6 +104,7 @@ std::vector<float> run_euler_sampler(
     // is the whole additive term, w_text * (v_cond - v_text) + w_speaker *
     // (v_cond - v_speaker), held in absolute units rather than as a ratio so a
     // reused correction cannot amplify a small v_cond.
+    const int n_cfg_lanes = options.cfg_mode == EchoCfgMode::Joint ? 2 : 3;
     std::vector<float> cfg_delta;
     int steps_since_refresh = 0;
 
@@ -98,12 +119,17 @@ std::vector<float> run_euler_sampler(
             // applied before one has been measured.
             const bool refresh = cfg_delta.empty() || steps_since_refresh >= cfg_interval - 1;
             if (refresh) {
-                auto lanes = denoise(x_t, t, 3);
-                if (static_cast<int64_t>(lanes.size()) != elements * 3) {
+                auto lanes = denoise(x_t, t, n_cfg_lanes);
+                if (static_cast<int64_t>(lanes.size()) != elements * n_cfg_lanes) {
                     throw std::runtime_error("Echo-TTS denoiser returned mis-shaped CFG lanes");
                 }
-                v_pred = combine_cfg_lanes(
-                    lanes, elements, options.cfg_scale_text, options.cfg_scale_speaker);
+                if (n_cfg_lanes == 2) {
+                    v_pred = combine_cfg_lanes_joint(
+                        lanes, elements, options.cfg_scale);
+                } else {
+                    v_pred = combine_cfg_lanes_independent(
+                        lanes, elements, options.cfg_scale_text, options.cfg_scale_speaker);
+                }
                 if (cfg_interval > 1) {
                     cfg_delta.resize(static_cast<size_t>(elements));
                     for (int64_t i = 0; i < elements; ++i) {