Function main

tests/ace_step/ace_step_warm_bench.cpp:243–348

}  // namespace

int main(int argc, char ** argv) {
    try {
        const std::filesystem::path model_path = arg_value(argc, argv, "--model", "models/Ace-Step1.5");
        const std::string backend_name = arg_value(argc, argv, "--backend", "cuda");
        const int device = int_arg(argc, argv, "--device", 0);
        const int threads = int_arg(argc, argv, "--threads", 1);
        const int warmup = int_arg(argc, argv, "--warmup", 0);
        const int iterations = int_arg(argc, argv, "--iterations", 1);
        const std::string request_sequence_json = arg_value(argc, argv, "--request-sequence-json", "");
        const std::string noise_file = arg_value(argc, argv, "--noise-file", "");
        const std::filesystem::path output_dir = arg_value(argc, argv, "--output-dir", "");
        const std::filesystem::path timing_path =
            arg_value(argc, argv, "--timing-file", "/tmp/ace_step_warm_bench_timing.log");
        if (request_sequence_json.empty()) {
            throw std::runtime_error("ACE-Step warmbench requires --request-sequence-json");
        }
        const auto session_option_overrides = parse_session_options(argc, argv);

        set_env_required("ENGINE_TRACE_ENABLED", "0");
        set_env_required("ENGINE_TIMING_ENABLED", "1");
        set_env_required("ENGINE_TIMING_FILE", timing_path.string());
        for (const auto & [key, value] : session_option_overrides) {
            if (key == "ace_step.dit_tensor_dump_dir") {
                set_env_required("ENGINE_ACE_STEP_DIT_TENSOR_DUMP_DIR", value);
            } else if (key == "ace_step.dit_tensor_compare_dir") {
                set_env_required("ENGINE_ACE_STEP_DIT_TENSOR_COMPARE_DIR", value);
            } else if (key == "ace_step.dit_tensor_compare_result") {
                set_env_required("ENGINE_ACE_STEP_DIT_TENSOR_COMPARE_RESULT", value);
            } else if (key == "ace_step.dit_tensor_compare_threshold") {
                set_env_required("ENGINE_ACE_STEP_DIT_TENSOR_COMPARE_THRESHOLD", value);
            } else if (key == "ace_step.trace_enabled") {
                set_env_required("ENGINE_TRACE_ENABLED", value);
            } else if (key == "ace_step.trace_file") {
                set_env_required("ENGINE_TRACE_FILE", value);
            }
        }
        engine::runtime::TaskSpec task{engine::runtime::VoiceTaskKind::Tts, engine::runtime::RunMode::Offline};
        auto registry = engine::runtime::make_default_registry();
        engine::runtime::ModelLoadRequest load_request;
        load_request.model_path = model_path;
        load_request.family_hint = "ace_step";
        for (const auto & [key, value] : session_option_overrides) {
            load_request.options.insert_or_assign(key, value);
        }
        auto model = registry.load(load_request);

        engine::runtime::SessionOptions options;
        options.backend.type = parse_backend(backend_name);
        options.backend.device = device;
        options.backend.threads = threads;
        for (const auto & [key, value] : session_option_overrides) {
            options.options.insert_or_assign(key, value);
        }
        auto session_base = model->create_task_session(task, options);
        auto * session = dynamic_cast<engine::runtime::IOfflineVoiceTaskSession *>(session_base.get());
        if (session == nullptr) {
            throw std::runtime_error("loaded ACE-Step session is not an offline TTS session");
        }