Update app.py

app.py

@@ -5,6 +5,9 @@ import jax
 import jax.numpy as jnp
 import gradio as gr
 
+from PIL import Image
+from transformers import ViTFeatureExtractor, AutoTokenizer, FlaxVisionEncoderDecoderModel
+
 from pathlib import Path
 from PIL import Image
 import numpy as np
@@ -61,12 +64,41 @@ pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
 
 
 
+loc = "ydshieh/vit-gpt2-coco-en"
+
+feature_extractor = ViTFeatureExtractor.from_pretrained(loc)
+tokenizer = AutoTokenizer.from_pretrained(loc)
+model = FlaxVisionEncoderDecoderModel.from_pretrained(loc)
+
+gen_kwargs = {"max_length": 16, "num_beams": 4}
+
+
+# This takes sometime when compiling the first time, but the subsequent inference will be much faster
+
+def generate(pixel_values):
+    output_ids = model.generate(pixel_values, **gen_kwargs).sequences
+    return output_ids
+    
+    
+def predict(image):
+
+    pixel_values = feature_extractor(images=image, return_tensors="np").pixel_values
+    output_ids = generate(pixel_values)
+    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+    preds = [pred.strip() for pred in preds]
+    
+    return preds
+    
+def image2text(image):
+    preds = predict(images[0])
+    return (preds[0])
+
 
 def text_to_image_and_image_to_text(text=None,image=None):
     txt=None
     img=None
     if image != None:
-        txt=text
+        txt=image2text(image)
     if text !=None:
         images = sd2_inference(pipeline, [text], params, seed = 42, num_inference_steps = 5 )
         img = images[0]