# 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 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. """ 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. """ # 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 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. """ 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 @staticmethod @spaces.GPU 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 # 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