Update audio_generator.py

audio_generator.py

@@ -1,34 +1,40 @@
+# audio_generation.py
+
+from transformers import AutoProcessor, BarkModel
+import torch
+import numpy as np
+from scipy.io.wavfile import write as write_wav
+from pydub import AudioSegment
 import os
-import edge_tts
-import asyncio
-import logging
-from datetime import datetime
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-async def _generate_speech(text: str, output_path: str) -> str:
-    try:
-        communicate = edge_tts.Communicate(
-            text=text,
-            voice="en-US-AriaNeural",
-            rate="+0%",
-            volume="+0%"
-        )
-        await communicate.save(output_path)
-        return output_path
-    except Exception as e:
-        logger.error(f"Generation failed: {str(e)}")
-        raise RuntimeError(f"Audio generation error: {str(e)}")
-
-def generate_audio(text: str) -> str:
-    """Main entry point for audio generation"""
-    try:
-        os.makedirs("tts_outputs", exist_ok=True)
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        output_path = os.path.join("tts_outputs", f"tts_{timestamp}.mp3")
-        
-        return asyncio.run(_generate_speech(text, output_path))
-    except Exception as e:
-        logger.error(f"Audio generation failed: {str(e)}")
-        raise
+import uuid
+
+processor = AutoProcessor.from_pretrained("suno/bark")
+model = BarkModel.from_pretrained("suno/bark")
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+def split_text(text, max_len=150):
+    return [text[i:i+max_len] for i in range(0, len(text), max_len)]
+
+def generate_audio(text, output_dir="audios"):
+    os.makedirs(output_dir, exist_ok=True)
+    chunks = split_text(text)
+    final_audio = AudioSegment.empty()
+
+    for idx, chunk in enumerate(chunks):
+        inputs = processor(chunk, return_tensors="pt").to(device)
+        audio_array = model.generate(**inputs)
+        audio_array = audio_array.cpu().numpy().squeeze()
+        audio_array = audio_array / np.max(np.abs(audio_array))
+
+        temp_path = os.path.join(output_dir, f"{uuid.uuid4()}.wav")
+        write_wav(temp_path, rate=22050, data=audio_array)
+
+        segment = AudioSegment.from_wav(temp_path)
+        final_audio += segment
+        os.remove(temp_path)
+
+    final_filename = os.path.join(output_dir, f"{uuid.uuid4()}_final.wav")
+    final_audio.export(final_filename, format="wav")
+    return final_filename