Update app.py

app.py

@@ -1,69 +1,57 @@
 import gradio as gr
-import torch
 from transformers import AutoProcessor, Blip2ForConditionalGeneration
+import torch
+from PIL import Image
 
-# Check if CUDA is available
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Load the BLIP-2 model and processor
+processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
+model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b")
 
-# Model ID
-MODEL_ID_FLAN_T5_XXL = "Salesforce/blip2-flan-t5-xxl"
+# Set device to GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
-# Load the model and processor
-processor = AutoProcessor.from_pretrained(MODEL_ID_FLAN_T5_XXL)
-model = Blip2ForConditionalGeneration.from_pretrained(MODEL_ID_FLAN_T5_XXL, load_in_8bit=True).to(device)
+def blip2_interface(image, prompted_caption_text, vqa_question, chat_context):
+    # Prepare image input
+    image_input = Image.fromarray(image).convert('RGB')
+    inputs = processor(image_input, return_tensors="pt").to(device, torch.float16)
+    
+    # Image Captioning
+    generated_ids = model.generate(**inputs, max_new_tokens=20)
+    image_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
 
-# Define a function for generating captions and answering questions
-def generate_text(image, text, decoding_method, temperature, length_penalty, repetition_penalty):
-    if text.startswith("Caption:"):
-        # Generate caption
-        inputs = processor(images=image, return_tensors="pt").to(device, torch.float16)
-        generated_ids = model.generate(
-            pixel_values=inputs.pixel_values,
-            do_sample=decoding_method == "Nucleus sampling",
-            temperature=temperature,
-            length_penalty=length_penalty,
-            repetition_penalty=repetition_penalty,
-            max_length=50,
-            min_length=1,
-            num_beams=5,
-            top_p=0.9,
-        )
-        result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
-        return result
-    else:
-        # Answer question
-        inputs = processor(images=image, text=text, return_tensors="pt").to(device, torch.float16)
-        generated_ids = model.generate(
-            **inputs,
-            do_sample=decoding_method == "Nucleus sampling",
-            temperature=temperature,
-            length_penalty=length_penalty,
-            repetition_penalty=repetition_penalty,
-            max_length=30,
-            min_length=1,
-            num_beams=5,
-            top_p=0.9,
-        )
-        result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
-        return result
+    # Prompted Image Captioning
+    inputs = processor(image_input, text=prompted_caption_text, return_tensors="pt").to(device, torch.float16)
+    generated_ids = model.generate(**inputs, max_new_tokens=20)
+    prompted_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+    
+    # Visual Question Answering (VQA)
+    prompt = f"Question: {vqa_question} Answer:"
+    inputs = processor(image_input, text=prompt, return_tensors="pt").to(device, torch.float16)
+    generated_ids = model.generate(**inputs, max_new_tokens=10)
+    vqa_answer = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+    
+    # Chat-based Prompting
+    prompt = chat_context + " Answer:"
+    inputs = processor(image_input, text=prompt, return_tensors="pt").to(device, torch.float16)
+    generated_ids = model.generate(**inputs, max_new_tokens=10)
+    chat_response = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+
+    return image_caption, prompted_caption, vqa_answer, chat_response
 
 # Define Gradio input and output components
-image_input = gr.Image(type="numpy")
-text_input = gr.Text()
+image_input = gr.inputs.Image(type="numpy")
+text_input = gr.inputs.Text()
 output_text = gr.outputs.Textbox()
 
-# Define Gradio interface
-gr.Interface(
-    fn=generate_text,
-    inputs=[image_input, text_input, gr.inputs.Radio(["Beam search", "Nucleus sampling"]), gr.inputs.Slider(0.5, 1.0, 0.1), gr.inputs.Slider(-1.0, 2.0, 0.2), gr.inputs.Slider(1.0, 5.0, 0.5)],
-    outputs=output_text,
-    examples=[
-        ["house.png", "Caption:"],
-        ["flower.jpg", "What is this flower and where is its origin?"],
-        ["pizza.jpg", "Caption:"],
-        ["sunset.jpg", "Caption:"],
-        ["forbidden_city.webp", "In what dynasties was this place built?"],
-    ],
-    title="BLIP-2",
-    description="Gradio demo for BLIP-2, image-to-text generation from Salesforce Research.",
-).launch()
+# Create Gradio interface
+iface = gr.Interface(
+    fn=blip2_interface,
+    inputs=[image_input, text_input, text_input, text_input],
+    outputs=[output_text, output_text, output_text, output_text],
+    title="BLIP-2 Image Captioning and VQA",
+    description="Interact with the BLIP-2 model for image captioning, prompted image captioning, visual question answering, and chat-based prompting.",
+)
+
+if __name__ == "__main__":
+    iface.launch()