Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
96 lines
3.3 KiB
C++
96 lines
3.3 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/assets/resource_bundle.h"
|
|
#include "engine/framework/runtime/session.h"
|
|
|
|
#include <cstdint>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
namespace engine::models::f5_tts {
|
|
|
|
struct F5RuntimeState;
|
|
|
|
// F5TTS_v1_Base / Habibi architecture (validated against the checkpoint:
|
|
// 22 blocks, dim 1024, 16 heads x 64, ff_mult 2, text_dim 512, conv_layers 4,
|
|
// 100 mel channels, embedding rows 2731).
|
|
struct F5Architecture {
|
|
int dim = 1024;
|
|
int depth = 22;
|
|
int heads = 16;
|
|
int head_dim = 64;
|
|
int ff_mult = 2;
|
|
int text_dim = 512;
|
|
int conv_layers = 4;
|
|
int mel_dim = 100;
|
|
int vocab_rows = 2731;
|
|
int sample_rate = 24000;
|
|
int hop_length = 256;
|
|
int n_fft = 1024;
|
|
};
|
|
|
|
struct F5SampleOptions {
|
|
float cfg_strength = 2.0F;
|
|
float sway_sampling_coef = -1.0F;
|
|
float speed = 1.0F;
|
|
int steps = 32;
|
|
uint32_t seed = 0;
|
|
};
|
|
|
|
// Compute device for the DiT forward: CUDA device index or CPU threads.
|
|
struct F5ComputeDevice {
|
|
bool use_cuda = false;
|
|
int device = 0; // CUDA device index
|
|
int threads = 0; // CPU threads; 0 = hardware concurrency
|
|
std::shared_ptr<F5RuntimeState> state = nullptr;
|
|
};
|
|
|
|
std::shared_ptr<F5RuntimeState> create_f5_lora_state(
|
|
const std::string & weights_path, const std::filesystem::path & adapter_path,
|
|
float strength, const F5ComputeDevice & device);
|
|
|
|
// Debug taps for parity testing: when non-null, intermediate stage outputs are
|
|
// appended (column layout, [features, T] flattened feature-major).
|
|
struct F5DebugTaps {
|
|
std::vector<float> * text_embed = nullptr; // after lookup + pos (01)
|
|
std::vector<float> * text_convnext = nullptr; // after 4 ConvNeXt (02)
|
|
std::vector<float> * text_padded = nullptr; // after pad/curtail (03)
|
|
std::vector<float> * input_embed = nullptr; // after proj + cpe (04)
|
|
std::vector<float> * time_embed = nullptr; // (05)
|
|
std::vector<float> * block0 = nullptr; // (07_block0)
|
|
std::vector<float> * block21 = nullptr; // (07_block21)
|
|
};
|
|
|
|
// Full DiT velocity forward for one Euler step inputs.
|
|
// x/cond: [seq_len * mel_dim] row-major (seq-major), text: ids (0 = filler),
|
|
// returns [seq_len * mel_dim] column layout (out[m * N + n]).
|
|
std::vector<float> f5_dit_forward(
|
|
const std::string & weights_path,
|
|
const std::vector<float> & x,
|
|
const std::vector<float> & cond,
|
|
const std::vector<int32_t> & text,
|
|
float time_value,
|
|
int seq_len,
|
|
const F5Architecture & arch,
|
|
bool drop_audio_cond,
|
|
bool drop_text,
|
|
const F5DebugTaps * taps = nullptr,
|
|
const F5ComputeDevice * device = nullptr);
|
|
|
|
// Batched CFG: one ne3=2 graph compute returning {conditioned, unconditioned}
|
|
// velocities. Matches python cfg_infer: the uncond half runs with
|
|
// drop_audio_cond (zeroed cond) and drop_text (filler text id 0); the host
|
|
// upload in runtime.cpp prepares both halves accordingly.
|
|
std::pair<std::vector<float>, std::vector<float>> f5_dit_forward_cfg(
|
|
const std::string & weights_path,
|
|
const std::vector<float> & x,
|
|
const std::vector<float> & cond,
|
|
const std::vector<int32_t> & text,
|
|
float time_value,
|
|
int seq_len,
|
|
const F5Architecture & arch,
|
|
const F5ComputeDevice * device = nullptr);
|
|
|
|
} // namespace engine::models::f5_tts
|