this is a total fucking mess.

app.py

@@ -1,57 +1,69 @@
 import gradio as gr
-from transformers import AutoProcessor, Blip2ForConditionalGeneration
 import torch
-from PIL import Image
-
-# 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")
+from transformers import AutoProcessor, Blip2ForConditionalGeneration
 
-# Set device to GPU if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+# Check if CUDA is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-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()
+# Model ID
+MODEL_ID_FLAN_T5_XXL = "Salesforce/blip2-flan-t5-xxl"
 
-    # 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()
+# 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)
 
-    return image_caption, prompted_caption, vqa_answer, chat_response
+# 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
 
 # Define Gradio input and output components
-image_input = gr.inputs.Image(type="numpy")
-text_input = gr.inputs.Text()
+image_input = gr.Image(type="numpy")
+text_input = gr.Text()
 output_text = gr.outputs.Textbox()
 
-# 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()
+# 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()