Update app.py

app.py

@@ -4,6 +4,7 @@ import os
 from huggingface_hub import list_liked_repos
 from diffusers import DiffusionPipeline
 import torch
+from transformers import AutoProcessor, Blip2ForConditionalGeneration
 
 def image_mod(prompt: str, model: str, image_0: gr.Image, image_1: gr.Image) -> list[gr.Image]:
     images = [image_0, image_1]
@@ -42,6 +43,29 @@ def change_model(choice):
     
     return images
 
+
+def give_dei_score(image_0: gr.Image, image_1: gr.Image):
+    textbox = []
+    
+    processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
+    model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device)
+
+    question = "Does this image present any racial or gender prejudices?"
+    prompt = f"Question: {question} Answer:"
+
+    images = [image_0, image_1]
+    for i, image in enumerate(images):
+        image_to_analyze = images[i]
+        inputs = processor(image_to_analyze, text=prompt, return_tensors="pt").to(device, torch.float16)
+        generated_ids = model.generate(**inputs, max_new_tokens=10)
+        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+        textbox.append(generated_text)
+    
+    return textbox
+
+
 with gr.Blocks() as demo:
     gr.Markdown(
     """
@@ -70,4 +94,12 @@ with gr.Blocks() as demo:
         image_mod, [textbox, dropdown, image_0, image_1], [image_0, image_1]
     )
 
+    
+    button = gr.Button()
+    image_0_textbox = gr.Textbox(label="Image 1")
+    image_1_textbox = gr.Textbox(label="Image 2")
+    button.click(
+        give_dei_score, [image_0, image_1], [image_0_textbox, image_1_textbox]
+    )
+
 demo.launch()