Update app.py

app.py

@@ -9,54 +9,98 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 model = OmniVoice.from_pretrained(
     "k2-fsa/OmniVoice",
-    # load_asr=True,
     device_map="cuda:0" if device == "cuda" else "cpu",
     dtype=torch.float16 if device == "cuda" else torch.float32
 )
 
-# model2 = OmniVoice.from_pretrained(
-#     "/root/.cache/huggingface/hub/models--k2-fsa--OmniVoice/snapshots/1d8c8a8fd2510535edab4f55aeae328b3e8a456e/",
-#     load_asr=True,
-#     asr_model_name="/root/.cache/huggingface/hub/models--openai--whisper-large-v3-turbo/snapshots/41f01f3fe87f28c78e2fbf8b568835947dd65ed9/",
-#     device_map="cuda:0" if device == "cuda" else "cpu",
-#     dtype=torch.float16 if device == "cuda" else torch.float32
-# )
+def build_voice_prompt(gender, age, pitch, style):
+    attrs = []
 
-def generate(text, ref_audio):
+    if gender:
+        attrs.append(gender)
+    if age:
+        attrs.append(age)
+    if pitch:
+        attrs.append(pitch)
+    if style:
+        attrs.append(style)
+
+    if len(attrs) > 0:
+        voice_desc = ", ".join(attrs)
+        return voice_desc
+    else:
+        return None
+
+
+def generate(text, ref_audio, gender, age, pitch, style, num_steps):
 
     output_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
 
+    instruct = build_voice_prompt(gender, age, pitch, style)
+
     if ref_audio is not None:
         audio = model.generate(
             text=text,
+            instruct=instruct,
             ref_audio=ref_audio,
-            num_step=4,
-            # ref_text optional
+            num_step=int(num_steps),
         )
     else:
         audio = model.generate(
             text=text,
-            num_step=4,
+            instruct=instruct,
+            num_step=int(num_steps),
         )
 
     sf.write(output_path, audio[0], 24000)
 
-    # paths = []
-    # for root, _, files in os.walk("/root/.cache/huggingface"):
-    #     for f in files:
-    #         paths.append(os.path.join(root, f))
-    # return "\n".join(paths)
-
     return output_path
 
+
 demo = gr.Interface(
     fn=generate,
     inputs=[
         gr.Textbox(label="Text"),
-        gr.Audio(type="filepath", label="Reference Voice (optional)")
+
+        gr.Audio(type="filepath", label="Reference Voice (optional)"),
+
+        gr.Dropdown(
+            choices=["male", "female"],
+            label="Gender",
+            value=None
+        ),
+
+        gr.Dropdown(
+            choices=["child", "teenager", "young adult", "middle-aged", "elderly"],
+            label="Age",
+            value=None
+        ),
+
+        gr.Dropdown(
+            choices=["very low pitch", "low pitch", "medium pitch", "high pitch"],
+            label="Pitch",
+            value=None
+        ),
+
+        # gr.Dropdown(
+        #     choices=["normal", "whisper", "calm", "angry"],
+        #     label="Style",
+        #     value=None
+        # ),
+
+        gr.Slider(
+            minimum=1,
+            maximum=32,
+            value=4,
+            step=1,
+            label="num_steps"
+        ),
     ],
+
     outputs=gr.Audio(type="filepath"),
-    title="OmniVoice TTS (Voice Cloning)"
+
+    title="OmniVoice TTS (Voice Design + Cloning)",
+    description="Control voice with gender, age, pitch, style + num_steps"
 )
 
 demo.launch()