Update app.py

app.py

@@ -1,47 +1,46 @@
-import gradio as gr
-from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
-import torch
-from PIL import Image
-import matplotlib.pyplot as plt
-
-# Load model and tokenizer
-model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
-# Define generation parameters
-max_length = 16
-num_beams = 4
-gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
-
-def predict_step(image):
-    # Process the input image
-    i_image = Image.fromarray(image)
-    if i_image.mode != "RGB":
-        i_image = i_image.convert(mode="RGB")
-    
-    # Prepare image for the model
-    pixel_values = feature_extractor(images=[i_image], return_tensors="pt").pixel_values
-    pixel_values = pixel_values.to(device)
-
-    # Generate prediction
-    output_ids = model.generate(pixel_values, **gen_kwargs)
-    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
-    preds = [pred.strip() for pred in preds]
-
-    return preds[0]
-
-# Define Gradio interface
-interface = gr.Interface(
-    fn=predict_step,
-    inputs=gr.inputs.Image(type="numpy"),
-    outputs="text",
-    title="Image-to-Text Conversion",
-    description="Upload an image, and this model will generate a textual description of the image."
-)
-
-# Launch the interface
-interface.launch()
+import gradio as gr
+from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
+import torch
+from PIL import Image
+
+# Load model and tokenizer
+model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# Define generation parameters
+max_length = 16
+num_beams = 4
+gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
+
+def predict_step(image):
+    # Convert image from numpy array to PIL image
+    i_image = Image.fromarray(image)
+    if i_image.mode != "RGB":
+        i_image = i_image.convert(mode="RGB")
+    
+    # Prepare image for the model
+    pixel_values = feature_extractor(images=[i_image], return_tensors="pt").pixel_values
+    pixel_values = pixel_values.to(device)
+
+    # Generate prediction
+    output_ids = model.generate(pixel_values, **gen_kwargs)
+    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+    preds = [pred.strip() for pred in preds]
+
+    return preds[0]
+
+# Define Gradio interface
+interface = gr.Interface(
+    fn=predict_step,
+    inputs=gr.Image(type="numpy"),  # Updated to use gr.Image
+    outputs="text",
+    title="Image-to-Text Conversion",
+    description="Upload an image, and this model will generate a textual description of the image."
+)
+
+# Launch the interface
+interface.launch()