diff --git a/include/s2_pipeline.h b/include/s2_pipeline.h index 19b3c3d..4ca3505 100644 --- a/include/s2_pipeline.h +++ b/include/s2_pipeline.h @@ -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; diff --git a/src/main.cpp b/src/main.cpp index d3e0b49..9d020d6 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -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 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"); @@ -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) { diff --git a/src/s2_codec.cpp b/src/s2_codec.cpp index e4dc3bc..e13654c 100755 --- a/src/s2_codec.cpp +++ b/src/s2_codec.cpp @@ -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); @@ -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 { @@ -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; } diff --git a/src/s2_pipeline.cpp b/src/s2_pipeline.cpp index 0dd5f01..1110c41 100644 --- a/src/s2_pipeline.cpp +++ b/src/s2_pipeline.cpp @@ -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( @@ -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);