Skip to content

feat: support separate codec GGUF via --codec-model #36

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jadsongmatos wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: support separate codec GGUF via --codec-model #36

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions include/s2_pipeline.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ class StreamingSink {

struct PipelineParams {
std::string model_path;
std::string codec_model_path;
std::string tokenizer_path;
std::string text;
std::string prompt_text;
Expand Down
6 changes: 4 additions & 2 deletions src/main.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,8 @@ void print_uso() {
safe_print(" --trim-silence Trim trailing silence in output WAV\n");
safe_print(" --no-normalize Keep original output peak level\n");
safe_print(" --normalize Peak-normalize output WAV to 0.95\n");
safe_print(" --codec-auto Benchmark codec backends and keep the fastest (default)\n");
safe_print(" --codec-model <path> Path to separate codec GGUF (split model files)\n");
safe_print(" --codec-auto Benchmark codec backends and keep the fastest (default)\n");
safe_print(" --codec-follow-backend Force codec to follow the selected GPU backend\n");
safe_print(" --codec-cpu Force codec on CPU even when model uses GPU\n");
safe_print(" --stream-file Write output WAV through the streaming path\n");
Expand Down Expand Up @@ -194,7 +195,8 @@ int main(int argc, char** argv) {
else if (arg == "--normalize") { params.normalize_output = true; }
else if (arg == "--codec-auto") { params.codec_auto_backend = true; params.codec_follow_backend = true; }
else if (arg == "--codec-follow-backend") { params.codec_auto_backend = false; params.codec_follow_backend = true; }
else if (arg == "--codec-cpu") { params.codec_auto_backend = false; params.codec_follow_backend = false; }
else if (arg == "--codec-cpu") { params.codec_auto_backend = false; params.codec_follow_backend = false; }
else if (arg == "--codec-model") { if (i+1 < argc) params.codec_model_path = argv[++i]; }
else if (arg == "--stream-file") { use_stream_file = true; }
else if (arg == "--stream-decode-stride") {
if (i+1 < argc) {
Expand Down
15 changes: 14 additions & 1 deletion src/s2_codec.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -747,7 +747,7 @@ bool AudioCodec::load_shared(SlowARModel* Model, gguf_context * shared_gguf_ctx,
return false;
}

if(!Model)
if(!Model || !shared_gguf_ctx)
{
struct gguf_init_params params = { true, &impl_->ctx_w };
gguf_context * local_gguf = gguf_init_from_file(gguf_path.c_str(), params);
Expand All @@ -762,6 +762,17 @@ bool AudioCodec::load_shared(SlowARModel* Model, gguf_context * shared_gguf_ctx,
else { impl_->ctx_w = Model->weights_.ctx_w; }

gguf_context * gguf_ctx = shared_gguf_ctx;
gguf_context * local_codec_gguf = nullptr;
if (!gguf_ctx && !gguf_path.empty()) {
struct gguf_init_params lparams = { true, nullptr };
local_codec_gguf = gguf_init_from_file(gguf_path.c_str(), lparams);
if (!local_codec_gguf) {
std::cerr << "[Codec] Failed to open GGUF for metadata: " << gguf_path << std::endl;
reset_codec_impl(*impl_);
return false;
}
gguf_ctx = local_codec_gguf;
}

try {
auto req_str = [&](const char * k) -> std::string {
Expand Down Expand Up @@ -880,9 +891,11 @@ bool AudioCodec::load_shared(SlowARModel* Model, gguf_context * shared_gguf_ctx,
impl_->residual_vq.clear();
} catch (const std::exception & e) {
std::cerr << "[Codec] " << e.what() << std::endl;
if (local_codec_gguf) gguf_free(local_codec_gguf);
reset_codec_impl(*impl_);
return false;
}
if (local_codec_gguf) gguf_free(local_codec_gguf);
S2_LOG_INFO_STREAM("[Codec] Backend: " << backend_name() << std::endl);
return true;
}
Expand Down
42 changes: 35 additions & 7 deletions src/s2_pipeline.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -477,7 +477,11 @@ bool Pipeline::init(const PipelineParams & params) {
safe_print_ln("Pipeline: loading codec on " + backend_label +
" device " + std::to_string(codec_gpu_device) + "...");
}
codec_loaded = codec().load_shared(&model(), shared_gguf, params.model_path, codec_gpu_device, codec_backend_type);
if (!params.codec_model_path.empty()) {
codec_loaded = codec().load_shared(&model(), nullptr, params.codec_model_path, codec_gpu_device, codec_backend_type);
} else {
codec_loaded = codec().load_shared(&model(), shared_gguf, params.model_path, codec_gpu_device, codec_backend_type);
}
if (!codec_loaded) {
if (codec_backend_type == BackendType::Metal) {
safe_print_warn_ln(
Expand All @@ -495,18 +499,42 @@ bool Pipeline::init(const PipelineParams & params) {
if (!use_gpu_codec) {
safe_print_ln("Pipeline: loading codec on CPU.");
}
codec_loaded = codec().load_shared(&model(), shared_gguf, params.model_path, -1, BackendType::CPU);
if (!params.codec_model_path.empty()) {
codec_loaded = codec().load_shared(&model(), nullptr, params.codec_model_path, -1, BackendType::CPU);
} else {
codec_loaded = codec().load_shared(&model(), shared_gguf, params.model_path, -1, BackendType::CPU);
}
}
if (!codec_loaded) {
safe_print_error_ln("Pipeline error: could not load codec from " + params.model_path);
std::string codec_src = params.codec_model_path.empty() ? params.model_path : params.codec_model_path;
safe_print_error_ln("Pipeline error: could not load codec from " + codec_src);
gguf_free(shared_gguf);
return false;
}

if (!read_all_tensor_data(params.model_path, shared_gguf, model(), codec())) {
safe_print_error_ln("Pipeline error: failed to read tensor data from " + params.model_path);
gguf_free(shared_gguf);
return false;
{
bool model_data_ok = read_all_tensor_data(params.model_path, shared_gguf, model(), codec());
if (!model_data_ok) {
safe_print_error_ln("Pipeline error: failed to read tensor data from " + params.model_path);
gguf_free(shared_gguf);
return false;
}
}

if (!params.codec_model_path.empty()) {
gguf_context * codec_gguf = gguf_init_from_file(params.codec_model_path.c_str(), gguf_params);
if (!codec_gguf) {
safe_print_error_ln("Pipeline error: failed to open codec GGUF from " + params.codec_model_path);
gguf_free(shared_gguf);
return false;
}
bool codec_data_ok = read_all_tensor_data(params.codec_model_path, codec_gguf, model(), codec());
gguf_free(codec_gguf);
if (!codec_data_ok) {
safe_print_error_ln("Pipeline error: failed to read codec tensor data from " + params.codec_model_path);
gguf_free(shared_gguf);
return false;
}
}

gguf_free(shared_gguf);
Expand Down