Update app.py

app.py

@@ -1,37 +1,22 @@
-import os
+import gradio as gr
 from PIL import Image
-import numpy as np
-import torch as T
-import transformers
+from transformers import pipeline
 
-# Assuming necessary model and tokenizer are already set up
-PATH_LLAVA = '_ckpt/LLaVA-7B-v1'
-tokenizer = transformers.AutoTokenizer.from_pretrained(PATH_LLAVA)
-model = transformers.AutoModelForCausalLM.from_pretrained(PATH_LLAVA).cuda()
+# Initialize the pipeline with the image captioning model
+caption_pipeline = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
 
-def remove_alter(s):  # Simplify expressive instruction
-    if 'ASSISTANT:' in s: s = s[s.index('ASSISTANT:')+10:].strip()
-    if '</s>' in s: s = s[:s.index('</s>')].strip()
-    if 'alternative' in s.lower(): s = s[:s.lower().index('alternative')]
-    if '[IMG0]' in s: s = s[:s.index('[IMG0]')]
-    s = '.'.join([s.strip() for s in s.split('.')[:2]])
-    if s[-1]!='.': s += '.'
-    return s.strip()
+def generate_caption(image):
+    # Convert the PIL Image to the format expected by the model
+    image = Image.open(image).convert("RGB")
+    
+    # Use the pipeline to generate a caption
+    result = caption_pipeline(image)
+    caption = result[0]["generated_text"]
+    
+    return caption
 
-def load_image_and_generate_instruction(image_path):
-    # Load the image
-    img = Image.open(image_path)
-    img.show()
-
-    # Example: Generate a simple instruction based on the image
-    # This is a placeholder. You would replace this with your own method
-    # to analyze the image and generate a textual description or instruction.
-    instruction = "Describe what to do with this image."
-
-    # Simplify and generate expressive instruction
-    expressive_instruction = remove_alter(instruction)
-    print("Expressive Instruction:", expressive_instruction)
-
-# Example usage
-image_path = './path/to/your/image.jpg'  # Update this path to your image
-load_image_and_generate_instruction(image_path)
+# Setup the Gradio interface
+interface = gr.Interface(fn=generate_caption,
+                         inputs=gr.inputs.Image(type="pil", label="Upload an Image"),
+                         outputs=gr.outputs.Textbox(label="Generated Caption"))
+interface.launch()