Update app.py

app.py

@@ -1,105 +1,70 @@
 import os
-import io
-import sys
-import time
-import gradio as gr
-from typing import List, Union, Any, Dict, Optional
 import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import PyPDF2
+import sounddevice as sd
 import numpy as np
-import librosa
-from transformers import SpeechSynthesisPipeline, AutoTokenizer, Wav2Vec2ForCTC, Wav2Vec2Processor
-from pypdfium2 import PdfDocument, Page
+from gtts import gTTS
+from io import BytesIO
+import gradio as gr
+
+def load_quantized_model(model_name):
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    # Quantize the model
+    model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
+    model.eval()
+
+    return model, tokenizer
 
-def convert_pdf_to_text(filepath):
-    doc = PdfDocument(filepath)
-    text = ""
-    for page in doc.pages():
-        text += page.extract_text() + "\n\n"
+def pdf_to_text(pdf_bytes):
+    pdf_file_obj = BytesIO(pdf_bytes)
+    pdf_reader = PyPDF2.PdfFileReader(pdf_file_obj)
+    text = ''
+    for page_num in range(pdf_reader.numPages):
+        page_obj = pdf_reader.getPage(page_num)
+        text += page_obj.extractText()
+    pdf_file_obj.close()
     return text
 
-class QuantizedSpeechT5TTSPipe:
-    def __init__(self):
-        self.tokenizer = AutoTokenizer.from_pretrained('facebook/wav2vec2-base-960h')
-        self.model = Wav2Vec2ForCTC.from_pretrained('/model/quantized_model').half().cuda()
-        self.processor = Wav2Vec2Processor.from_pretrained("/model/quantized_vocab")
-        
-    def _pad_and_concatenate(self, tensor_list: List[torch.Tensor], padding_value=0):
-        max_length = max([tensor.size(0) for tensor in tensor_list])
-        padded_tensors = []
-        for tensor in tensor_list:
-            pad_width = max_length - tensor.size(0)
-            if pad_width > 0:
-                tensor_padded = torch.cat((tensor, torch.full((pad_width, tensor.size(1)), fill_value=padding_value).type_as(tensor)))
-            else:
-                tensor_padded = tensor
-            padded_tensors.append(tensor_padded)
-        
-        return torch.stack(padded_tensors)
-    
-    def preprocess(self, inputs: Union[str, List[str]], **kwargs) -> dict:
-        if isinstance(inputs, str):
-            inputs = [inputs]
-        batch_encodings = self.tokenizer(inputs, truncation=True, padding='longest', return_tensors="pt").input_values
-        return {"batch_encodings": batch_encodings}
-    
-    def postprocess(self, outputs: Dict[str, torch.Tensor], **kwargs) -> Union[List[str], List[bytes]]:
-        logits = outputs["logits"].cpu().detach().numpy()
-        ids = np.argmax(logits, axis=-1)
-        cleaned_ids = [id_seq[:np.where(np.array(id_seq) == 2)[0][0]] for id_seq in ids] # Remove CTC blanks
-        decoded_strings = self.tokenizer.decode(cleaned_ids)
-        
-        audios = []
-        for text in decoded_strings:
-            input_values = self.processor(text, sampling_rate=16000, return_tensors="pt").input_values
-            input_values = input_values.cuda().unsqueeze(0)
-            
-            mel_outputs = self.model(input_values).mel_output
-            _, predicted_ids = torch.topk(mel_outputs.float(), k=1, dim=-1)
-            predicted_ids = predicted_ids.squeeze(-1).tolist()[0]
-            
-            raw_waveform = self.processor.post_processing(predicted_ids)
-            waveform = raw_waveform * 32768 / max(abs(raw_waveform))
-            wav_data = np.int16(waveform)
-            audio = io.BytesIO()
-            sf.write(audio, int(44100), wav_data, format="WAV")
-            audios.append(audio.getvalue())
-        return audios
-    
-    def generate(self, text: str):
-        processed_inputs = self.preprocess(text)
-        outputs = self.model(**processed_inputs)
-        results = self.postprocess(outputs)
-        return results
+def generate_audio(model, tokenizer, text):
+    input_ids = torch.tensor(tokenizer.encode(text, return_tensors="pt")).cuda()
+    with torch.no_grad():
+        outputs = model.generate(input_ids, max_length=500, pad_token_id=tokenizer.eos_token_id)
+    output_text = tokenizer.decode(outputs[0])
+    return output_text
 
-if __name__ == "__main__":
-    tts = QuantizedSpeechT5TTSPipe()
-    sample_text = 'Hello world! This is a test.'
-    result = tts.generate(sample_text)
-    print(f'Generated {len(result)} audio files from "{sample_text}"')
+def save_and_play_audio(text):
+    tts = gTTS(text=text, lang='en')
+    output_file = "output.mp3"
+    tts.save(output_file)
+
+    data, fs = sd.default.read_audio(output_file)
+    sd.play(data, fs)
+    sd.wait()
+
+    return output_file
 
-def main(pdf_file: gr.File, output_filename: str):
-    start_time = time.time()
-    pdf_text = convert_pdf_to_text(pdf_file)
-    print(f'Processed PDF content in {time.time() - start_time:.4f} seconds')
+def main(pdf_file):
+    # Load the quantized model
+    model, tokenizer = load_quantized_model("microsoft/speecht5_tts")
 
-    pipe = QuantizedSpeechT5TTSPipe()
-    start_time = time.time()
-    audios = pipe.generate(pdf_text)
-    print(f'Generated {len(audios)} audio files in {time.time() - start_time:.4f} seconds')
+    # Move the model to the GPU if available
+    if torch.cuda.is_available():
+        model.cuda()
 
-    zip_buffer = BytesIO()
-    with ZipFile(zip_buffer, mode='w') as zf:
-        for i, audio in enumerate(audios):
-            filename = f"{i}_{output_filename}.wav"
-            zf.writestr(filename, audio)
-    zip_buffer.seek(0)
+    # Convert the uploaded PDF file to text
+    text = pdf_to_text(pdf_file.read())
 
-    return {'zip': zip_buffer}
+    # Generate audio from the text
+    audio_text = generate_audio(model, tokenizer, text)
 
-iface = gr.Interface(fn=main, 
-                      inputs="file", 
-                      outputs="binary",
-                      input_types=['pdf'],
-                      output_types=['download'])
+    # Save and play the audio
+    output_file = save_and_play_audio(audio_text)
 
-iface.launch()
+    return {"output_file": output_file}
+
+if __name__ == "__main__":
+    app = gr.Interface(main, inputs=gr.inputs.File(type="pdf"), outputs="text")
+    app.launch()