Update app.py

app.py

@@ -1,24 +1,28 @@
 import gradio as gr
 from transformers import pipeline
-# from TTS.api import TTS
 
 import librosa
 import numpy as np
 import torch
 
 from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
+from transformers import AutoProcessor, AutoModelForCausalLM
 
 
 checkpoint = "microsoft/speecht5_tts"
-processor = SpeechT5Processor.from_pretrained(checkpoint)
-model = SpeechT5ForTextToSpeech.from_pretrained(checkpoint)
+tts_processor = SpeechT5Processor.from_pretrained(checkpoint)
+tts_model = SpeechT5ForTextToSpeech.from_pretrained(checkpoint)
 vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
 
+
+vqa_processor = AutoProcessor.from_pretrained("microsoft/git-base-textvqa")
+vqa_model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-textvqa")
+
 def tts(text):
     if len(text.strip()) == 0:
         return (16000, np.zeros(0).astype(np.int16))
 
-    inputs = processor(text=text, return_tensors="pt")
+    inputs = tts_processor(text=text, return_tensors="pt")
 
     # limit input length
     input_ids = inputs["input_ids"]
@@ -44,26 +48,36 @@ def tts(text):
 
     speaker_embedding = torch.tensor(speaker_embedding).unsqueeze(0)
 
-    speech = model.generate_speech(input_ids, speaker_embedding, vocoder=vocoder)
+    speech = tts_model.generate_speech(input_ids, speaker_embedding, vocoder=vocoder)
 
     speech = (speech.numpy() * 32767).astype(np.int16)
     return (16000, speech)
 
 
-captioner = pipeline(model="microsoft/git-base")
+# captioner = pipeline(model="microsoft/git-base")
 # tts = TTS(model_name="tts_models/multilingual/multi-dataset/your_tts", progress_bar=False, gpu=False)
 
 
 def predict(image):
-    text = captioner(image)[0]["generated_text"]
+    # text = captioner(image)[0]["generated_text"]
 
     # audio_output = "output.wav"
     # tts.tts_to_file(text, speaker=tts.speakers[0], language="en", file_path=audio_output)
+
+    pixel_values = vqa_processor(images=image, return_tensors="pt").pixel_values
+
+    prompt = "what is in the scene?"
+    prompt_ids = vqa_processor(text=prompt, add_special_tokens=False).input_ids
+    prompt_ids = [vqa_processor.tokenizer.cls_token_id] + prompt_ids
+    prompt_ids = torch.tensor(prompt_ids).unsqueeze(0)
+    
+    text_ids = vqa_model.generate(pixel_values=pixel_values, input_ids=prompt_ids, max_length=50)
+    text = vqa_processor.batch_decode(text_ids, skip_special_tokens=True)
+    
     audio = tts(text)
     
     return text, audio
 
-# theme = gr.themes.Default(primary_hue="#002A5B")
 
 demo = gr.Interface(
     fn=predict,
@@ -74,6 +88,3 @@ demo = gr.Interface(
 )
 
 demo.launch()
-
-# gr.Interface.load("models/ronniet/git-base-env").launch()
-# gr.Interface.load("models/microsoft/git-base").launch()