Update app.py

app.py

@@ -7,13 +7,17 @@ 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")
 
+ic_processor = AutoProcessor.from_pretrained("microsoft/git-base")
+ic_model = AutoModelForCausalLM.from_pretrained("microsoft/git-base")
+
 def tts(text):
     if len(text.strip()) == 0:
         return (16000, np.zeros(0).astype(np.int16))
@@ -22,7 +26,7 @@ def tts(text):
 
     # limit input length
     input_ids = inputs["input_ids"]
-    input_ids = input_ids[..., :model.config.max_text_positions]
+    input_ids = input_ids[..., :tts_model.config.max_text_positions]
 
     # if speaker == "Surprise Me!":
     #     # load one of the provided speaker embeddings at random
@@ -44,21 +48,26 @@ 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 = ic_processor(images=image, return_tensors="pt").pixel_values
+    text_ids = ic_model.generate(pixel_values=pixel_values, max_length=50)
+    text = ic_processor.batch_decode(text_ids, skip_special_tokens=True)[0]
+    
     audio = tts(text)
     
     return text, audio
@@ -74,6 +83,3 @@ demo = gr.Interface(
 )
 
 demo.launch()
-
-# gr.Interface.load("models/ronniet/git-base-env").launch()
-# gr.Interface.load("models/microsoft/git-base").launch()