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
71 lines
3 KiB
C++
71 lines
3 KiB
C++
#pragma once
|
|
|
|
#include "engine/community_models/f5_tts/runtime.h"
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/core/backend.h"
|
|
#include "engine/framework/core/backend_weight_store.h"
|
|
#include "engine/framework/core/module.h"
|
|
#include "engine/framework/modules/activation_modules.h"
|
|
#include "engine/framework/modules/attention/scaled_dot_product_attention.h"
|
|
#include "engine/framework/modules/conv_modules.h"
|
|
#include "engine/framework/modules/linear_module.h"
|
|
#include "engine/framework/modules/lookup_modules.h"
|
|
#include "engine/framework/modules/norm_modules.h"
|
|
#include "engine/framework/modules/positional_modules.h"
|
|
#include "engine/framework/modules/primitive_modules.h"
|
|
#include "engine/framework/modules/streaming_conv_modules.h"
|
|
#include "engine/framework/modules/structural_modules.h"
|
|
|
|
#include <cstdint>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace engine::models::f5_tts {
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Weights, expressed with framework module weight types (dev-branch pattern:
|
|
// the weight store loads torch-shaped tensors once; modules consume them).
|
|
// ---------------------------------------------------------------------------
|
|
|
|
struct F5TextConvNextWeights {
|
|
modules::DepthwiseConv1dWeights dwconv; // [C, 1, 7] + [C]
|
|
modules::NormWeights norm; // [C] + [C]
|
|
modules::LinearWeights pw1; // [1024, 512]
|
|
modules::LinearWeights pw2; // [512, 1024]
|
|
std::vector<float> grn_gamma; // [1024] (host: GRN has no module)
|
|
std::vector<float> grn_beta; // [1024]
|
|
};
|
|
|
|
struct F5BlockWeights {
|
|
modules::LinearWeights attn_norm; // adaLN: [6*D, D]
|
|
modules::LinearWeights to_q; // [D, D]
|
|
modules::LinearWeights to_k;
|
|
modules::LinearWeights to_v;
|
|
modules::LinearWeights to_out;
|
|
modules::LinearWeights ff0; // [2*D, D]
|
|
modules::LinearWeights ff2; // [D, 2*D]
|
|
};
|
|
|
|
struct F5DiTWeights {
|
|
int64_t vocab_size = 2731; // from the checkpoint's text embedding
|
|
std::shared_ptr<core::BackendWeightStore> store;
|
|
core::TensorValue text_embedding; // [vocab, 512]
|
|
modules::LinearWeights input_proj; // [1024, 712]
|
|
modules::Conv1dWeights cpe0; // [1024, 64, 31] (grouped g=16)
|
|
modules::Conv1dWeights cpe2;
|
|
modules::LinearWeights time0; // [256, 256]
|
|
modules::LinearWeights time2; // [1024, 256]
|
|
std::vector<F5TextConvNextWeights> text_blocks; // x4
|
|
std::vector<F5BlockWeights> blocks; // x22
|
|
modules::LinearWeights norm_out; // [2*D, D]
|
|
modules::LinearWeights proj_out; // [100, D]
|
|
};
|
|
|
|
// Loads the EMA weights (prefix stripped) into module weight types.
|
|
F5DiTWeights load_dit_weights(
|
|
const assets::TensorSource & source,
|
|
ggml_backend_t backend,
|
|
core::BackendType backend_type);
|
|
|
|
} // namespace engine::models::f5_tts
|