Updates

app.py

@@ -1,62 +1,35 @@
 import gradio as gr
 from PIL import Image
-import onnxruntime as ort
+import torch
 import torchvision.transforms as transforms
 import json
 import os
 import numpy as np
 import pandas as pd
 import random
-from huggingface_hub import snapshot_download, HfApi
-from transformers import CLIPTokenizer
+import onnxruntime as ort
+from huggingface_hub import HfApi
+from transformers import CLIPTokenizer, AutoImageProcessor, AutoModelForImageClassification
+from safetensors.torch import load_file as safe_load
 
 # --- Config ---
 HUB_REPO_ID = "CDL-AMLRT/OpenArenaLeaderboard"
 HF_TOKEN = os.environ.get("HF_TOKEN")
 LOCAL_JSON = "leaderboard.json"
 HUB_JSON = "leaderboard.json"
-MODEL_PATH = "mobilenet_v2_fake_detector.onnx"
+MODEL_PATH = "model.safetensors"  # ✅ updated filename
+MODEL_BACKBONE = "microsoft/swinv2-small-patch4-window16-256"
 CLIP_IMAGE_ENCODER_PATH = "clip_image_encoder.onnx"
 CLIP_TEXT_ENCODER_PATH = "clip_text_encoder.onnx"
 PROMPT_CSV_PATH = "generate2_1.csv"
 PROMPT_MATCH_THRESHOLD = 10  # percent
 
-# --- Download leaderboard + model checkpoint from HF Hub ---
+# --- No-op for HF Space ---
 def load_assets():
-    try:
-        snapshot_download(
-            repo_id=HUB_REPO_ID,
-            local_dir=".",
-            repo_type="dataset",
-            token=HF_TOKEN,
-            allow_patterns=[HUB_JSON, MODEL_PATH, CLIP_IMAGE_ENCODER_PATH, CLIP_TEXT_ENCODER_PATH, PROMPT_CSV_PATH]
-        )
-    except Exception as e:
-        print(f"Failed to load assets from HF Hub: {e}")
+    print("Skipping snapshot_download. Assuming files exist via Git LFS in HF Space.")
 
 load_assets()
 
-# --- Load prompts from CSV ---
-def load_prompts():
-    try:
-        df = pd.read_csv(PROMPT_CSV_PATH)
-        if "prompt" in df.columns:
-            return df["prompt"].dropna().tolist()
-        else:
-            print("CSV missing 'prompt' column.")
-            return []
-    except Exception as e:
-        print(f"Failed to load prompts: {e}")
-        return []
-
-PROMPT_LIST = load_prompts()
-
-def load_initial_state():
-    sorted_scores = sorted(leaderboard_scores.items(), key=lambda x: x[1], reverse=True)
-    leaderboard_table = [[name, points] for name, points in sorted_scores]
-    return gr.update(value=get_random_prompt()), leaderboard_table
-
-
 # --- Load leaderboard ---
 def load_leaderboard():
     try:
@@ -66,10 +39,8 @@ def load_leaderboard():
         print(f"Failed to read leaderboard: {e}")
         return {}
 
-
 leaderboard_scores = load_leaderboard()
 
-# --- Save and push to HF Hub ---
 def save_leaderboard():
     try:
         with open(HUB_JSON, "w", encoding="utf-8") as f:
@@ -91,10 +62,33 @@ def save_leaderboard():
     except Exception as e:
         print(f"Failed to save leaderboard to HF Hub: {e}")
 
-# --- Load ONNX models ---
-session = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
-input_name = session.get_inputs()[0].name
+# --- Load prompts from CSV ---
+def load_prompts():
+    try:
+        df = pd.read_csv(PROMPT_CSV_PATH)
+        if "prompt" in df.columns:
+            return df["prompt"].dropna().tolist()
+        else:
+            print("CSV missing 'prompt' column.")
+            return []
+    except Exception as e:
+        print(f"Failed to load prompts: {e}")
+        return []
+
+PROMPT_LIST = load_prompts()
+
+# --- Load model + processor ---
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+processor = AutoImageProcessor.from_pretrained(MODEL_BACKBONE)
+model = AutoModelForImageClassification.from_pretrained(MODEL_BACKBONE)
+model.classifier = torch.nn.Linear(model.config.hidden_size, 2)
 
+model.load_state_dict(safe_load(MODEL_PATH, device="cpu"), strict=False)
+model.to(device)
+model.eval()
+
+# --- CLIP prompt matching ---
 clip_image_sess = ort.InferenceSession(CLIP_IMAGE_ENCODER_PATH, providers=["CPUExecutionProvider"])
 clip_text_sess = ort.InferenceSession(CLIP_TEXT_ENCODER_PATH, providers=["CPUExecutionProvider"])
 clip_tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
@@ -134,19 +128,22 @@ def detect_with_model(image: Image.Image, prompt: str, username: str):
     prompt_score = compute_prompt_match(image, prompt)
     if prompt_score < PROMPT_MATCH_THRESHOLD:
         message = f"⚠️ Prompt match too low ({round(prompt_score, 2)}%). Please generate an image that better matches the prompt."
-        return message, None, [], gr.update(visible=True), gr.update(visible=False), username
+        return message, None, leaderboard, gr.update(visible=True), gr.update(visible=False), username
 
-    image_tensor = transforms.Resize((224, 224))(image)
-    image_tensor = transforms.ToTensor()(image_tensor).unsqueeze(0).numpy().astype(np.float32)
-    outputs = session.run(None, {input_name: image_tensor})
-    prob = round(1 / (1 + np.exp(-outputs[0][0][0])), 2)
-    prediction = "Real" if prob > 0.5 else "Fake"
+    # Run model inference
+    inputs = processor(image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+        pred_class = torch.argmax(logits, dim=-1).item()
+        prediction = "Real" if pred_class == 0 else "Fake"
 
-    score = 1 if prediction == "Real" else 0
-    confidence = round(prob * 100, 2) if prediction == "Real" else round((1 - prob) * 100, 2)
+        probs = torch.softmax(logits, dim=-1)[0]
+        confidence = round(probs[pred_class].item() * 100, 2)
 
-    message = f"🔍 Prediction: {prediction} ({round(confidence, 2)}% confidence)\n🧐 Prompt match: {prompt_score}%"
+    score = 1 if prediction == "Real" else 0
 
+    message = f"🔍 Prediction: {prediction} ({confidence}% confidence)\n🧐 Prompt match: {round(prompt_score, 2)}%"
     if prediction == "Real":
         leaderboard_scores[username] = leaderboard_scores.get(username, 0) + score
         message += "\n🎉 Nice! You fooled the AI. +1 point!"
@@ -167,17 +164,22 @@ def detect_with_model(image: Image.Image, prompt: str, username: str):
         username
     )
 
-# --- UI Layout ---
 def get_random_prompt():
     return random.choice(PROMPT_LIST) if PROMPT_LIST else "A synthetic scene with dramatic lighting"
 
+def load_initial_state():
+    sorted_scores = sorted(leaderboard_scores.items(), key=lambda x: x[1], reverse=True)
+    leaderboard_table = [[name, points] for name, points in sorted_scores]
+    return gr.update(value=get_random_prompt()), leaderboard_table
+
+# --- Gradio UI ---
 with gr.Blocks(css=".gr-button {font-size: 16px !important}") as demo:
     gr.Markdown("## 🌝 OpenFake Arena")
     gr.Markdown("Welcome to the OpenFake Arena!\n\n**Your mission:** Generate a synthetic image for the prompt, upload it, and try to fool the AI detector into thinking it’s real.\n\n**Rules:**\n- Only synthetic images allowed!\n- No cheating with real photos.\n\nMake it wild. Make it weird. Most of all — make it fun.")
 
     with gr.Group(visible=True) as input_section:
         username_input = gr.Textbox(label="Your Name", placeholder="Enter your name", interactive=True)
-        model_input = gr.Textbox(label="Model Used", placeholder="Name of tge model used to generate the image", interactive=True)
+        model_input = gr.Textbox(label="Model Used", placeholder="Name of the model used to generate the image", interactive=True)
 
         with gr.Row():
             prompt_input = gr.Textbox(
@@ -225,7 +227,16 @@ with gr.Blocks(css=".gr-button {font-size: 16px !important}") as demo:
     )
 
     try_again_btn.click(
-        fn=lambda name: ("", None, [], gr.update(visible=True), gr.update(visible=False), name, gr.update(value=get_random_prompt())),
+        fn=lambda name: (
+            "",               # Clear prediction text
+            None,             # Clear uploaded image
+            leaderboard,               # Clear leaderboard (temporarily, gets reloaded on next submit)
+            gr.update(visible=True),   # Show input section
+            gr.update(visible=False),  # Hide "Try Again" button
+            name,             # Keep username
+            gr.update(value=get_random_prompt()),  # Load new prompt
+            None              # Clear image input
+        ),
         inputs=[username_input],
         outputs=[
             prediction_output,
@@ -234,14 +245,16 @@ with gr.Blocks(css=".gr-button {font-size: 16px !important}") as demo:
             input_section,
             try_again_btn,
             username_input,
-            prompt_input
+            prompt_input,
+            image_input        # ← added output to clear image
         ]
     )
 
+
     demo.load(
         fn=load_initial_state,
         outputs=[prompt_input, leaderboard]
-    )   
+    )
 
 if __name__ == "__main__":
     demo.launch()