Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import os
 import gradio as gr
 import torch
@@ -13,11 +12,11 @@ DEFAULT_PROMPT = "<image>"
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-model = FlamingoModel.from_pretrained('dhansmair/flamingo-tiny')
+model = FlamingoModel.from_pretrained('dhansmair/flamingo-mini')
 model.to(device)
 model.eval()
 
-processor = FlamingoProcessor(model.config, load_vision_processor=True)
+processor = FlamingoProcessor(model.config, device=device)
 
 # setup some example images
 examples = []
@@ -30,10 +29,11 @@ if os.path.isdir(EXAMPLES_DIR):
 def predict_caption(image, prompt):
     assert isinstance(prompt, str)
     
-    features = processor.extract_features(image).to(device)
-    caption = model.generate_captions(processor, 
-            visual_features=features,
-            prompt=prompt)
+    caption = model.generate_captions(
+        processor, 
+        images=image,
+        prompt=prompt
+    )
   
     if isinstance(caption, list):
         caption = caption[0]
@@ -46,4 +46,4 @@ iface = gr.Interface(fn=predict_caption,
         examples=examples,
         outputs="text")
 
-iface.launch()
+iface.launch()