Allow setting seed (PABannier#105)

harperreed · Sep 6, 2023 · 5e2826d · 5e2826d
1 parent 79b2cbe
commit 5e2826d
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 0 deletions.
diff --git a/bark.cpp b/bark.cpp
@@ -184,6 +184,12 @@ struct bark_context * bark_new_context_with_model(struct bark_model * model) {
  return ctx;
 }
 
+void bark_seed_rng(struct bark_context * ctx, int32_t seed) {
+ if (ctx) {
+ ctx->rng.seed(seed);
+ }
+}
+
 int bark_vocab_load(
  const char * fname,
  bark_vocab * vocab,

diff --git a/bark.h b/bark.h
@@ -66,6 +66,8 @@ extern "C" {
 
  BARK_API struct bark_context * bark_new_context_with_model(struct bark_model * model);
 
+ BARK_API void bark_seed_rng(struct bark_context * ctx, int32_t seed);
+
  BARK_API void bark_free(struct bark_context * ctx);
 
  BARK_API void bark_free_model(struct bark_model * ctx);

diff --git a/examples/main.cpp b/examples/main.cpp
@@ -100,6 +100,8 @@ int main(int argc, char **argv) {
 
  printf("\n");
 
+ bark_seed_rng(bctx, params.seed);
+
  const int64_t t_eval_us_start = ggml_time_us();
  bark_generate_audio(bctx, params.prompt.data(), params.dest_wav_path.c_str(), params.n_threads);
  t_eval_us = ggml_time_us() - t_eval_us_start;