Create app.py

app.py

@@ -0,0 +1,120 @@
+import os
+import torch
+import gradio as gr
+from PIL import Image
+from lavis.models import load_model_and_preprocess, model_zoo
+
+# Fix CUDA plugin registration errors
+os.environ["CUDA_VISIBLE_DEVICES"] = "0" if torch.cuda.is_available() else "-1"
+
+class InstructBLIP:
+    def __init__(self):
+        self.model = None
+        self.vis_processors = None
+        self.txt_processors = None
+        self.device = "cpu"
+
+    def load_models(self, model, vis_processors, txt_processors, device):
+        self.model = model
+        self.vis_processors = vis_processors
+        self.txt_processors = txt_processors
+        self.device = device
+
+    def query(self, image, question):
+        image = self.vis_processors["eval"](image).unsqueeze(0).to(self.device)
+        samples = {"image": image, "prompt": question}
+        candidates = ["yes", "no"]
+        ans = self.model.predict_class(samples=samples, candidates=candidates)
+
+        # Convert logits to probabilities
+        logits = ans[0]
+        yes_prob = torch.sigmoid(torch.tensor(logits[0])).item() * 100
+        no_prob = torch.sigmoid(torch.tensor(logits[1])).item() * 100
+
+        result = "Real" if yes_prob > no_prob else "Fake"
+        confidence = max(yes_prob, no_prob)
+
+        return result, round(confidence, 2)
+
+def load_model(model_name="blip2_t5", model_type="pretrain_flant5xl"):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+
+    try:
+        model, vis_processors, txt_processors = load_model_and_preprocess(
+            name=model_name, 
+            model_type=model_type, 
+            is_eval=True, 
+            device=device
+        )
+        if model is None:
+            raise ValueError(f"Failed to load model '{model_name}' with type '{model_type}'")
+
+        instruct = InstructBLIP()
+        instruct.load_models(model, vis_processors, txt_processors, device)
+        return instruct
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None
+
+
+# Load the model once when the script starts
+model_instance = load_model()
+
+def predict_image(input_image, question="Is this photo real [*]?"):
+    if input_image is None:
+        return "No image provided", 0
+
+    try:
+        # Ensure input is a PIL Image
+        if not isinstance(input_image, Image.Image):
+            input_image = Image.fromarray(input_image)
+
+        # Run model inference
+        result, confidence = model_instance.query(input_image, question)
+        return result, confidence
+    except Exception as e:
+        return f"Error: {str(e)}", 0
+
+# Create Gradio interface
+def create_interface():
+    with gr.Blocks(title="Fake Image Detector") as app:
+        gr.Markdown("""
+        # Real vs Fake Image Detector
+        Upload an image to check if it's real or AI-generated. The model will classify the image and provide a confidence score.
+        Based on AntifakePrompt: https://github.com/nctu-eva-lab/AntifakePrompt
+        """)
+
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(type="pil", label="Upload Image")
+                question = gr.Textbox(label="Question Prompt", value="Is this photo real [*]?")
+                submit_btn = gr.Button("Analyze Image", variant="primary")
+
+            with gr.Column():
+                result_label = gr.Textbox(label="Classification Result")
+                confidence = gr.Number(label="Confidence Score (%)")
+
+        submit_btn.click(
+            fn=predict_image,
+            inputs=[input_image, question],
+            outputs=[result_label, confidence]
+        )
+
+        gr.Examples(
+            examples=[
+                ["example_real.jpg", "Is this photo real [*]?"],
+                ["example_fake.jpg", "Is this photo real [*]?"],
+            ],
+            inputs=[input_image, question],
+            outputs=[result_label, confidence],
+            fn=predict_image,
+            cache_examples=True,
+        )
+
+    return app
+
+if __name__ == "__main__":
+    # Create and launch the Gradio interface
+    demo = create_interface()
+    demo.launch(share=True)  # Set share=True to get a public link