k

app.py

@@ -1,33 +1,27 @@
 import gradio as gr
-from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
-import torch
-from PIL import Image
+from transformers import pipeline
 
-model_name = "nlpconnect/vit-gpt2-image-captioning"
-model = VisionEncoderDecoderModel.from_pretrained(model_name)
-feature_extractor = ViTImageProcessor.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+# Load pipeline
+captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
 
+# Define Gradio interface function
 def generate_caption(image):
-    if image is None:
-        return "Please upload an image."
-
-    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
-    pixel_values = pixel_values.to(device)
+    result = captioner(image)
+    return result[0]["generated_text"].strip()
 
-    output_ids = model.generate(pixel_values, max_length=16, num_beams=4)
-    caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return caption
+# Create Gradio app
+title = "🖼️ Image Caption Generator using ViT-GPT2"
+description = "Upload an image and get a descriptive caption using the `nlpconnect/vit-gpt2-image-captioning` model by Hugging Face Transformers."
 
 demo = gr.Interface(
     fn=generate_caption,
     inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="🖼️ Image Caption Generator",
-    description="Upload an image and get a caption describing it using a VisionEncoderDecoder model (ViT + GPT-2)."
+    outputs=gr.Textbox(label="Generated Caption"),
+    title=title,
+    description=description,
+    allow_flagging="never"
 )
 
-demo.launch()
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -1,4 +1,3 @@
-transformers
+transformers>=4.37.0
 torch
-gradio
-Pillow
+gradio>=4.0.0