Upload app.py

app.py

@@ -0,0 +1,48 @@
+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()