app.py

import spaces
import gradio as gr
import torch
from transformers import MarianTokenizer, MarianMTModel
from parler_tts import ParlerTTSForConditionalGeneration
from transformers import AutoTokenizer, AutoFeatureExtractor, set_seed
from PyPDF2 import PdfReader
import re
import textwrap
import soundfile as sf

# Device configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Initialize models and tokenizers
tts_model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-large-v1").to(device)
tts_tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-large-v1")
feature_extractor = AutoFeatureExtractor.from_pretrained("parler-tts/parler-tts-mini-v1")
SAMPLE_RATE = feature_extractor.sampling_rate
SEED = 42

# Helper function to extract text from a PDF
def pdf_to_text(pdf_path):
    with open(pdf_path, 'rb') as file:
        pdf_reader = PdfReader(file)
        text = ""
        for page_num in range(len(pdf_reader.pages)):
            page = pdf_reader.pages[page_num]
            text += page.extract_text()
    return text

# Helper function to split text into sentences using regex
def split_text_into_sentences(text):
    sentence_endings = re.compile(r'[.!?]')
    sentences = sentence_endings.split(text)
    return [sentence.strip() for sentence in sentences if sentence.strip()]

# Translation function
@spaces.GPU(duration=120)
def translate(source_text, source_lang, target_lang, batch_size=16):
    model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
    
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name).to(device)
    
    text_chunks = textwrap.wrap(source_text, 512)
    translated_text = ""
    
    for i in range(0, len(text_chunks), batch_size):
        text_batch = text_chunks[i:i+batch_size]
        input_ids = tokenizer(text_batch, return_tensors="pt", padding=True, truncation=True, max_length=512).input_ids.to(device)
        output_ids = model.generate(input_ids, max_new_tokens=512)
        
        for output in output_ids:
            output_text = tokenizer.decode(output, skip_special_tokens=True)
            translated_text += output_text + " "
    
    return translated_text

# Function to preprocess the text (normalization, punctuation)
def preprocess(text):
    text = text.replace("-", " ")
    if text[-1] not in ".!?":
        text += "."
    return text

# Function to generate audio for a single sentence
@spaces.GPU(duration=120)
def generate_single_wav_from_text(sentence, description):
    set_seed(SEED)
    inputs = tts_tokenizer(description.strip(), return_tensors="pt").to(device)
    prompt = tts_tokenizer(preprocess(sentence), return_tensors="pt").to(device)

    generation = tts_model.generate(
        input_ids=inputs.input_ids, prompt_input_ids=prompt.input_ids, attention_mask=inputs.attention_mask,
        prompt_attention_mask=prompt.attention_mask, do_sample=True, temperature=1.0
    )
    audio_arr = generation.cpu().numpy().squeeze()
    return SAMPLE_RATE, audio_arr


# Gradio Interface
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            pdf_input = gr.File(label="Upload PDF", file_types=['pdf'])
            translate_checkbox = gr.Checkbox(label="Enable Translation", value=False)
            source_lang = gr.Dropdown(choices=["en", "tr", "de", "fr"], label="Source Language", value="en", interactive=True)
            target_lang = gr.Dropdown(choices=["tr"], label="Target Language", value="tr", interactive=True)
            description = gr.Textbox(label="Voice Description", lines=2, 
                                     value="Old man voice. Monotone voice tune from an old man, with a very close recording that almost has no background noise.")
            run_button = gr.Button("Generate Audio", variant="primary")
        with gr.Column():
            audio_gallery = gr.Gallery(label="Generated Audios", item_type="audio")
            markdown_output = gr.Markdown()

    def handle_process(pdf_input, translate_checkbox, source_lang, target_lang, description):
        text = pdf_to_text(pdf_input.name)
        if translate_checkbox:
            text = translate(text, source_lang, target_lang)
        
        sentences = split_text_into_sentences(text)
        all_audio_paths = []
        all_text = ""

        for sentence in sentences:
            sample_rate, audio_arr = generate_single_wav_from_text(sentence, description)
            # Create temporary audio file
            with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as tmpfile:
                sf.write(tmpfile, audio_arr, sample_rate)
                all_audio_paths.append(tmpfile.name)
            all_text += f"**Sentence**: {sentence}\n\n"
            yield all_audio_paths, all_text

    def run_pipeline(pdf_input, translate_checkbox, source_lang, target_lang, description):
        for audio_data, markdown_text in handle_process(pdf_input, translate_checkbox, source_lang, target_lang, description):
            yield audio_data, markdown_text

    translate_checkbox.change(fn=handle_translation_toggle, inputs=translate_checkbox, outputs=[source_lang, target_lang])
    source_lang.change(fn=lambda lang: gr.update(choices={"en": ["de", "fr", "tr"], "tr": ["en"], "de": ["en", "fr"], "fr": ["en", "de"]}.get(lang, [])), inputs=source_lang, outputs=target_lang)
    run_button.click(run_pipeline, inputs=[pdf_input, translate_checkbox, source_lang, target_lang, description], outputs=[audio_gallery, markdown_output])

demo.queue()
demo.launch(share=True)