update audio

app.py

@@ -44,6 +44,32 @@ def transcribe(ref_audio, language=None):
         return_timestamps=False,
     )["text"].strip()
 
+MODES = {
+    "Student Only (4 steps)": {
+        "teacher_steps": 0,
+        "teacher_stopping_time": 1.0,
+        "student_start_step": 0,
+        "description": "Fastest (4 steps), good quality"
+    },
+    "Teacher-Guided (8 steps)": {
+        "teacher_steps": 16,
+        "teacher_stopping_time": 0.07,
+        "student_start_step": 1,
+        "description": "Best balance (8 steps), recommended"
+    },
+    "High Diversity (16 steps)": {
+        "teacher_steps": 24,
+        "teacher_stopping_time": 0.3,
+        "student_start_step": 2,
+        "description": "More natural prosody (16 steps)"
+    },
+    "Custom": {
+        "teacher_steps": None,
+        "teacher_stopping_time": None,
+        "student_start_step": None,
+        "description": "Fine-tune all parameters"
+    }
+}
 
 @spaces.GPU(duration=120)
 def generate_speech(
@@ -67,24 +93,11 @@ def generate_speech(
         prompt_text = transcribe(prompt_audio)
 
 
-    if mode == "Student Only (4 steps)":
-        teacher_steps = 0
-        student_start_step = 0
-        teacher_stopping_time = 1.0
-    elif mode == "Teacher-Guided (8 steps)":
-        teacher_steps = 16
-        teacher_stopping_time = 0.07
-        student_start_step = 1
-    elif mode == "High Diversity (16 steps)":
-        teacher_steps = 24
-        teacher_stopping_time = 0.3
-        student_start_step = 2
-    else:  # Custom
-        teacher_steps = custom_teacher_steps
-        teacher_stopping_time = custom_teacher_stopping_time
-        student_start_step = custom_student_start_step
-
-    # Generate speech
+    if mode == "Custom":
+        teacher_steps, teacher_stopping_time, student_start_step = custom_teacher_steps, custom_teacher_stopping_time, custom_student_start_step
+    else:
+        teacher_steps, teacher_stopping_time, student_start_step = MODES[mode].values()
+
     generated_audio = model.generate(
         gen_text=target_text,
         audio_path=prompt_audio,
@@ -97,27 +110,15 @@ def generate_speech(
     )
 
 
-    # Save audio
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
-        output_path = tmp_file.name
-
-    if isinstance(generated_audio, np.ndarray):
-        generated_audio = torch.from_numpy(generated_audio)
-
-    if generated_audio.dim() == 1:
-        generated_audio = generated_audio.unsqueeze(0)
+    if isinstance(generated_audio, torch.Tensor):
+        audio_np = generated_audio.cpu().numpy()
+    else:
+        audio_np = generated_audio
 
-    torchaudio.save(output_path, generated_audio, 24000)
+    if audio_np.ndim == 1:
+        audio_np = np.expand_dims(audio_np, axis=0)
 
-    return (
-        output_path,
-        "Success!",
-        (
-            f"Mode: {mode} | Transcribed: {prompt_text[:50]}..."
-            if not prompt_text
-            else f"Mode: {mode}"
-        ),
-    )
+    return (24000, audio_np)
 
 
 # Create Gradio interface