Update app.py

app.py

@@ -2,23 +2,22 @@ import os
 import gradio as gr
 import torch
 import PIL
-
-from flamingo_mini import FlamingoConfig, FlamingoModel, FlamingoProcessor
-
-
+from transformers import AutoProcessor, AutoModelForCausalLM  # Using AutoModel classes
 
 EXAMPLES_DIR = 'examples'
 DEFAULT_PROMPT = "<image>"
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-model = FlamingoModel.from_pretrained('dhansmair/flamingo-mini')
+# Load model using AutoModel with trust_remote_code=True
+model = AutoModelForCausalLM.from_pretrained('dhansmair/flamingo-mini', trust_remote_code=True)
 model.to(device)
 model.eval()
 
-processor = FlamingoProcessor(model.config, device=device)
+# Initialize processor without the `device` argument
+processor = AutoProcessor.from_pretrained('dhansmair/flamingo-mini')
 
-# setup some example images
+# Setup some example images
 examples = []
 if os.path.isdir(EXAMPLES_DIR):
     for file in os.listdir(EXAMPLES_DIR):
@@ -29,10 +28,10 @@ if os.path.isdir(EXAMPLES_DIR):
 def predict_caption(image, prompt):
     assert isinstance(prompt, str)
     
-    caption = model.generate_captions(
-        processor, 
-        images=image,
-        prompt=prompt
+    # Process the image using the model
+    caption = model.generate(
+        processor(images=image, prompt=prompt),  # Pass processed inputs to the model
+        max_length=50
     )
   
     if isinstance(caption, list):
@@ -41,9 +40,11 @@ def predict_caption(image, prompt):
     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 = gr.Interface(
+    fn=predict_caption, 
+    inputs=[gr.Image(type="pil"), gr.Textbox(value=DEFAULT_PROMPT, label="Prompt")], 
+    examples=examples,
+    outputs="text"
+)
 
 iface.launch(debug=True)