upd: audio output

app.py

@@ -5,8 +5,14 @@ from transformers import (
     AutoProcessor,
     AutoModelForDocumentQuestionAnswering,
     pipeline,
+    VitsModel,
 )
 import torch
+import numpy as np
+
+
+mms_tts_model = VitsModel.from_pretrained("facebook/mms-tts-rus")
+mms_tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-rus")
 
 processor = AutoProcessor.from_pretrained(
     "MariaK/layoutlmv2-base-uncased_finetuned_docvqa_v2"
@@ -71,6 +77,17 @@ def generate_answer(image, question):
     return answer_ru
 
 
+def text_to_speech(text):
+
+    inputs = mms_tts_tokenizer(text, return_tensors="pt")
+
+    with torch.no_grad():
+        output = mms_tts_model(**inputs).waveform
+
+    audio = output.numpy()
+    return text, (16000, audio.squeeze())
+
+
 def transcribe(image, audio):
     if not image or not audio:
         return
@@ -79,6 +96,10 @@ def transcribe(image, audio):
 
     if y.ndim > 1:
         y = y.mean(axis=1)
+
+    y = y.astype(np.float32)
+    y /= np.max(np.abs(y))
+
     transcription_text = transcriber({"sampling_rate": sr, "raw": y})["text"]
 
     return generate_answer(image, transcription_text)
@@ -90,7 +111,10 @@ qa_interface = gr.Interface(
         gr.Image(type="pil"),
         gr.Textbox(label="Вопрос (на русском)", placeholder="Ваш вопрос"),
     ],
-    outputs=gr.Textbox(label="Ответ (на русском)"),
+    outputs=[
+        gr.Textbox(label="Ответ (на русском)"),
+        gr.Audio(label="Сгенерированное аудио"),
+    ],
     examples=[["doc.png", "О чем данный документ?"]],
     live=False,
 )
@@ -101,7 +125,10 @@ speech_interface = gr.Interface(
         gr.Image(type="pil"),
         gr.Audio(sources="microphone", label="Голосовой ввод"),
     ],
-    outputs=gr.Textbox(label="Распознанный текст"),
+    outputs=[
+        gr.Textbox(label="Ответ (на русском)"),
+        gr.Audio(label="Сгенерированное аудио"),
+    ],
     live=True,
 )
 interface = gr.TabbedInterface(