generate_audio.py

# generate_audio.py

import spaces
import pickle
import torch
import numpy as np
from tqdm import tqdm
from transformers import BarkModel, AutoProcessor, AutoTokenizer
from parler_tts import ParlerTTSForConditionalGeneration
from scipy.io import wavfile
from pydub import AudioSegment
import io
import ast


@spaces.GPU
class TTSGenerator:
    """
    A class to generate podcast-style audio from a transcript using ParlerTTS and Bark models.
    """
    #@spaces.GPU
    def __init__(self, transcript_file_path,output_audio_path):
        """
        Initialize the TTS generator with the path to the rewritten transcript file.
        
        Args:
            transcript_file_path (str): Path to the file containing the rewritten transcript.
        """
        self.transcript_file_path = transcript_file_path
        self.output_audio_path = output_audio_path
        
        # Set device
        self.device = "cuda" if torch.cuda.is_available() else "cpu"


        # Load Parler model and tokenizer for Speaker 1
        self.parler_model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-mini-v1").to(self.device)
        self.parler_tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-mini-v1")
        self.speaker1_description = """
        Laura's voice is expressive and dramatic in delivery, speaking at a moderately fast pace with a very close recording that almost has no background noise and very clear audio.
        """
        self.speaker2_description = """
        Gary's voice is expressive and dramatic in delivery, speaking at a moderately fast pace with a very close recording that almost has no background noise and very clear audio.
        """
        
        # Load Bark model and processor for Speaker 2
        # self.bark_processor = AutoProcessor.from_pretrained("suno/bark-small")
        # self.bark_model = BarkModel.from_pretrained("suno/bark-small", torch_dtype=torch.float16).to(self.device)
        # self.bark_sampling_rate = 24000
        # self.voice_preset = "v2/en_speaker_6"

    #@spaces.GPU
    def load_transcript(self):
        """
        Loads the rewritten transcript from the specified file.
        
        Returns:
            list: The content of the transcript as a list of tuples (speaker, text).
        """
        with open(self.transcript_file_path, 'rb') as f:
            return ast.literal_eval(pickle.load(f))
    
    #@spaces.GPU(duration=30)
    def generate_speaker1_audio(self, text):
        """
        Generate audio for Speaker 1 using ParlerTTS.
        
        Args:
            text (str): Text to be synthesized for Speaker 1.
        
        Returns:
            np.array: Audio array.
            int: Sampling rate.
        """
        # input_ids = self.parler_tokenizer(self.speaker1_description, return_tensors="pt").input_ids.to(self.device)
        # prompt_input_ids = self.parler_tokenizer(text, return_tensors="pt").input_ids.to(self.device)
        # generation = self.parler_model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids)
        # audio_arr = generation.cpu().numpy().squeeze()
        # return audio_arr, self.parler_model.config.sampling_rate
        with torch.no_grad():
            input_ids = self.parler_tokenizer(self.speaker1_description, return_tensors="pt", padding=True).input_ids.to(self.device)
            attention_mask_input = self.parler_tokenizer(self.speaker1_description, return_tensors="pt", padding=True).attention_mask.to(self.device)
            
            prompt_input_ids = self.parler_tokenizer(text, return_tensors="pt", padding=True).input_ids.to(self.device)
            attention_mask_prompt = self.parler_tokenizer(text, return_tensors="pt", padding=True).attention_mask.to(self.device)
            
            # Pass all required arguments to generate() for reliable behavior
            generation = self.parler_model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask_input,  # Set attention mask for input IDs
                prompt_input_ids=prompt_input_ids,
                prompt_attention_mask=attention_mask_prompt  # Set prompt attention mask
            )
            audio_arr = generation.cpu().numpy().squeeze()
        return audio_arr, self.parler_model.config.sampling_rate

    #@spaces.GPU(duration=30)
    def generate_speaker2_audio(self, text):
        """
        Generate audio for Speaker 2 using Bark.
        
        Args:
            text (str): Text to be synthesized for Speaker 2.
        
        Returns:
            np.array: Audio array.
            int: Sampling rate.
        """
        with torch.no_grad():
            input_ids = self.parler_tokenizer(self.speaker2_description, return_tensors="pt", padding=True).input_ids.to(self.device)
            attention_mask_input = self.parler_tokenizer(self.speaker1_description, return_tensors="pt", padding=True).attention_mask.to(self.device)
            
            prompt_input_ids = self.parler_tokenizer(text, return_tensors="pt", padding=True).input_ids.to(self.device)
            attention_mask_prompt = self.parler_tokenizer(text, return_tensors="pt", padding=True).attention_mask.to(self.device)
            
            # Pass all required arguments to generate() for reliable behavior
            generation = self.parler_model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask_input,  # Set attention mask for input IDs
                prompt_input_ids=prompt_input_ids,
                prompt_attention_mask=attention_mask_prompt  # Set prompt attention mask
            )
            audio_arr = generation.cpu().numpy().squeeze()
        
        # inputs = self.bark_processor(text, voice_preset=self.voice_preset).to(self.device)
        # speech_output = self.bark_model.generate(**inputs, temperature=0.9, semantic_temperature=0.8)
        # audio_arr = speech_output[0].cpu().numpy()
        # return audio_arr, self.bark_sampling_rate
        # Tokenize input text and obtain input IDs and attention mask
        # inputs = self.bark_processor(text, voice_preset=self.voice_preset).to(self.device)
        # speech_output = self.bark_model.generate(**inputs, temperature=0.9, semantic_temperature=0.8)
        # audio_arr = speech_output[0].cpu().numpy()
        return audio_arr, self.parler_model.config.sampling_rate

    
    #@spaces.GPU
    @staticmethod
    def numpy_to_audio_segment(audio_arr, sampling_rate):
        """
        Convert numpy array to AudioSegment.
        
        Args:
            audio_arr (np.array): Numpy array of audio data.
            sampling_rate (int): Sampling rate of the audio.
        
        Returns:
            AudioSegment: Converted audio segment.
        """
        audio_int16 = (audio_arr * 32767).astype(np.int16)
        byte_io = io.BytesIO()
        wavfile.write(byte_io, sampling_rate, audio_int16)
        byte_io.seek(0)
        return AudioSegment.from_wav(byte_io)
    
    #@spaces.GPU(duration=300)
    def generate_audio(self):
        """
        Converts the transcript into audio and saves it to a file.
        
        Returns:
            str: Path to the saved audio file.
        """
        transcript = self.load_transcript()
        final_audio = None

        for speaker, text in tqdm(transcript, desc="Generating podcast segments", unit="segment"):
            if speaker == "Speaker 1":
                audio_arr, rate = self.generate_speaker1_audio(text)
            else:  # Speaker 2
                audio_arr, rate = self.generate_speaker2_audio(text)
            
            # Convert to AudioSegment
            audio_segment = self.numpy_to_audio_segment(audio_arr, rate)
            
            # Add segment to final audio
            if final_audio is None:
                final_audio = audio_segment
            else:
                final_audio += audio_segment
            torch.cuda.empty_cache()

        # Export final audio to MP3
        final_audio.export(self.output_audio_path, format="mp3", bitrate="192k", parameters=["-q:a", "0"])
        return self.output_audio_path