app.py

import os
import gradio as gr
import torch
import PIL

from flamingo_mini import FlamingoConfig, FlamingoModel, FlamingoProcessor

EXAMPLES_DIR = 'examples'
DEFAULT_PROMPT = "<image>"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = FlamingoModel.from_pretrained('dhansmair/flamingo-mini')
model.to(device)
model.eval()
processor = FlamingoProcessor(model.config, load_vision_processor=True)

# setup some example images
examples = []
if os.path.isdir(EXAMPLES_DIR):
    for file in os.listdir(EXAMPLES_DIR):
        path = EXAMPLES_DIR + "/" + file
        examples.append([path, DEFAULT_PROMPT])

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)
    if isinstance(caption, list):
        caption = caption[0]
    return caption
    
iface = gr.Interface(fn=predict_caption, 
        inputs=[gr.Image(type="pil"), gr.Textbox(value=DEFAULT_PROMPT, label="Prompt")], 
        examples=examples,
        outputs="text")

iface.launch()