add more prompts

hy133 · Apr 11, 2023 · 7c88255 · 7c88255
1 parent 20990e2
commit 7c88255
Show file tree

Hide file tree

Showing 18 changed files with 8 additions and 9 deletions.
diff --git a/README.md b/README.md
@@ -80,7 +80,7 @@ text_prompt = """
     I have a silky smooth voice, and today I will tell you about 
     the exercise regimen of the common sloth.
 """
-audio_array = generate_audio(text_prompt, history_prompt="man-narrator")
+audio_array = generate_audio(text_prompt, history_prompt="speech_0")
 ```
 
 [sloth.webm](https://user-images.githubusercontent.com/5068315/230684883-a344c619-a560-4ff5-8b99-b4463a34487b.webm)

diff --git a/bark/assets/prompts/brylcream.npz b/bark/assets/prompts/brylcream.npz
diff --git a/bark/assets/prompts/es-woman.npz b/bark/assets/prompts/es-woman.npz
diff --git a/bark/assets/prompts/music_0.npz b/bark/assets/prompts/music_0.npz
diff --git a/bark/assets/prompts/music_1.npz b/bark/assets/prompts/music_1.npz
diff --git a/bark/assets/prompts/music_2.npz b/bark/assets/prompts/music_2.npz
diff --git a/bark/assets/prompts/music_3.npz b/bark/assets/prompts/music_3.npz
diff --git a/bark/assets/prompts/music_4.npz b/bark/assets/prompts/music_4.npz
diff --git a/bark/assets/prompts/music_5.npz b/bark/assets/prompts/music_5.npz
diff --git a/bark/assets/prompts/man-narrator.npz → bark/assets/prompts/speech_0.npz b/bark/assets/prompts/man-narrator.npz → bark/assets/prompts/speech_0.npz
diff --git a/bark/assets/prompts/speech_1.npz b/bark/assets/prompts/speech_1.npz
diff --git a/bark/assets/prompts/speech_2.npz b/bark/assets/prompts/speech_2.npz
diff --git a/bark/assets/prompts/speech_3.npz b/bark/assets/prompts/speech_3.npz
diff --git a/bark/assets/prompts/speech_4.npz b/bark/assets/prompts/speech_4.npz
diff --git a/bark/assets/prompts/speech_5.npz b/bark/assets/prompts/speech_5.npz
diff --git a/bark/assets/prompts/speech_6.npz b/bark/assets/prompts/speech_6.npz
diff --git a/bark/assets/prompts/speech_7.npz b/bark/assets/prompts/speech_7.npz
diff --git a/bark/generation.py b/bark/generation.py
@@ -49,10 +49,9 @@ def autocast():
 SAMPLE_RATE = 24_000
 
 
-ALLOWED_PROMPTS = (
-    "brylcream",
-    "es-woman",
-    "man-narrator",
+ALLOWED_PROMPTS = set(
+    [f"speech_{n}" for n in range(8)] +
+    [f"music_{n}" for n in range(6)]
 )
 
 
@@ -307,7 +306,7 @@ def generate_text_semantic(
         assert (history_prompt in ALLOWED_PROMPTS)
         semantic_history = np.load(
             os.path.join(CUR_PATH, "assets", "prompts", f"{history_prompt}.npz")
-        )["text"]
+        )["semantic_prompt"]
         assert (
             isinstance(semantic_history, np.ndarray)
             and len(semantic_history.shape) == 1
@@ -452,8 +451,8 @@ def generate_coarse(
         x_history = np.load(
             os.path.join(CUR_PATH, "assets", "prompts", f"{history_prompt}.npz")
         )
-        x_semantic_history = x_history["coarse_1"]
-        x_coarse_history = x_history["coarse_2"]
+        x_semantic_history = x_history["semantic_prompt"]
+        x_coarse_history = x_history["coarse_prompt"]
         assert (
             isinstance(x_semantic_history, np.ndarray)
             and len(x_semantic_history.shape) == 1
@@ -594,7 +593,7 @@ def generate_fine(
         assert (history_prompt in ALLOWED_PROMPTS)
         x_fine_history = np.load(
             os.path.join(CUR_PATH, "assets", "prompts", f"{history_prompt}.npz")
-        )["fine"]
+        )["fine_prompt"]
         assert (
             isinstance(x_fine_history, np.ndarray)
             and len(x_fine_history.shape) == 2