app.py

app.py

@@ -0,0 +1,41 @@
+import gradio as gr
+import torch
+from transformers import GPT2TokenizerFast, ViTImageProcessor, VisionEncoderDecoderModel 
+
+# Setup device, model, tokenizer, and feature extractor
+device = 'cpu'
+model_checkpoint = "Stoneman/IG-caption-generator-nlpconnect-all"
+feature_extractor = ViTImageProcessor.from_pretrained(model_checkpoint)
+tokenizer = GPT2TokenizerFast.from_pretrained(model_checkpoint)
+model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(device)
+
+# Prediction function
+def predict(image, max_length=128):
+    image = image.convert('RGB')
+    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
+    caption_ids = model.generate(pixel_values, max_length=max_length)[0]
+    caption_text = tokenizer.decode(caption_ids, skip_special_tokens=True)
+    return caption_text
+
+# Define input and output components
+input_component = gr.components.Image(label="Upload any Image", type="pil")
+output_component = gr.components.Textbox(label="Captions")
+
+# Example images
+examples = [f"example{i}.JPG" for i in range(1, 10)]
+
+# Interface
+title = "IG-caption-generator"
+description = "Made by: Jiayu Shi"
+interface = gr.Interface(
+    fn=predict,
+    description=description,
+    inputs=input_component,
+    theme="huggingface",
+    outputs=output_component,
+    examples=examples,
+    title=title,
+)
+
+# Launch interface
+interface.launch(debug=True)

requirements.txt

@@ -0,0 +1,2 @@
+transformers
+torch