Update app.py

app.py

@@ -21,9 +21,6 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print("=" * 60)
 print("FORESIGHT - MODEL INITIALIZATION")
 print("=" * 60)
-print(f"Model repository : {MODEL_REPO}")
-print(f"Device           : {DEVICE}")
-print(f"HF token present : {HF_TOKEN is not None}")
 
 processor = ViTImageProcessor(
     do_resize=True,
@@ -35,32 +32,15 @@ processor = ViTImageProcessor(
     image_std=[0.229, 0.224, 0.225],
 )
 
-print("Loading custom model code...")
+config_path = hf_hub_download(repo_id=MODEL_REPO, filename="configuration_vit_refiner.py", token=HF_TOKEN)
+model_path = hf_hub_download(repo_id=MODEL_REPO, filename="modeling_vit_refiner.py", token=HF_TOKEN)
 
-# Download configuration and modeling scripts from Hugging Face Hub
-config_path = hf_hub_download(
-    repo_id=MODEL_REPO,
-    filename="configuration_vit_refiner.py",
-    token=HF_TOKEN,
-)
-model_path = hf_hub_download(
-    repo_id=MODEL_REPO,
-    filename="modeling_vit_refiner.py",
-    token=HF_TOKEN,
-)
-
-# Dynamically import configuration_vit_refiner
-config_spec = importlib.util.spec_from_file_location(
-    "configuration_vit_refiner", config_path
-)
+config_spec = importlib.util.spec_from_file_location("configuration_vit_refiner", config_path)
 config_module = importlib.util.module_from_spec(config_spec)
 sys.modules["configuration_vit_refiner"] = config_module
 config_spec.loader.exec_module(config_module)
 
-# Dynamically import modeling_vit_refiner
-model_spec = importlib.util.spec_from_file_location(
-    "modeling_vit_refiner", model_path
-)
+model_spec = importlib.util.spec_from_file_location("modeling_vit_refiner", model_path)
 model_module = importlib.util.module_from_spec(model_spec)
 sys.modules["modeling_vit_refiner"] = model_module
 model_spec.loader.exec_module(model_module)
@@ -68,32 +48,18 @@ model_spec.loader.exec_module(model_module)
 ViTRefinerConfig = config_module.ViTRefinerConfig
 ViTRefinerModel = model_module.ViTRefinerModel
 
-# Instantiate model directly to prevent nested meta-context issues
 config = ViTRefinerConfig.from_pretrained(MODEL_REPO, token=HF_TOKEN)
 model = ViTRefinerModel(config)
 
-# Load weights safely from safetensors
-weights_path = hf_hub_download(
-    repo_id=MODEL_REPO,
-    filename="model.safetensors",
-    token=HF_TOKEN,
-)
+weights_path = hf_hub_download(repo_id=MODEL_REPO, filename="model.safetensors", token=HF_TOKEN)
 state_dict = load_file(weights_path)
 model.load_state_dict(state_dict, strict=False)
 
 model.to(DEVICE)
 model.eval()
-print("Model initialized successfully.")
-
-# Override labels to REAL / FAKE
-model.config.id2label = {
-    0: "FAKE",
-    1: "REAL"
-}
-model.config.label2id = {
-    "FAKE": 0,
-    "REAL": 1
-}
+
+model.config.id2label = {0: "FAKE", 1: "REAL"}
+model.config.label2id = {"FAKE": 0, "REAL": 1}
 
 # -----------------------------
 # ATTENTION HOOK EXTRACTION
@@ -102,30 +68,37 @@ model.config.label2id = {
 captured_attentions = []
 
 def attention_hook(module, input, output):
-    # Standard ViTSelfAttention outputs (context_layer, attention_probs)
-    if isinstance(output, tuple) and len(output) > 1:
-        captured_attentions.append(output[1].detach().cpu())
-
-# Attach hooks to all internal ViT self-attention layers
-for name, module in model.vit.vit.named_modules():
-    if "attention.self" in name or name.endswith("self"):
+    # If output is a tuple with attention probabilities, grab index 1
+    if isinstance(output, tuple):
+        if len(output) > 1 and output[1] is not None:
+            captured_attentions.append(output[1].detach().cpu())
+        elif isinstance(output[0], torch.Tensor) and output[0].ndim == 4: # [batch, heads, tokens, tokens]
+            captured_attentions.append(output[0].detach().cpu())
+
+# Register hook on every self-attention block
+for module in model.modules():
+    if module.__class__.__name__ == "ViTSelfAttention":
         module.register_forward_hook(attention_hook)
 
 # -----------------------------
-# ATTENTION ROLLOUT (Exact from reference)
+# ATTENTION ROLLOUT
 # -----------------------------
 
 def compute_attention_rollout(attentions):
     if not attentions:
         return None
 
-    # Shape: [layers, batch, heads, tokens, tokens]
-    att_mat = torch.stack(attentions).squeeze(1)
+    # Stack along layers dimension -> [layers, batch, heads, tokens, tokens]
+    att_mat = torch.stack(attentions)
+    if att_mat.dim() == 5:
+        att_mat = att_mat[:, 0] # Select batch item 0 -> [layers, heads, tokens, tokens]
+
+    # Mean over attention heads -> [layers, tokens, tokens]
     att_mat = att_mat.mean(dim=1)
 
     residual_att = torch.eye(att_mat.size(-1), device=att_mat.device)
     aug_att_mat = att_mat + residual_att
-    aug_att_mat = aug_att_mat / aug_att_mat.sum(dim=-1).unsqueeze(-1)
+    aug_att_mat = aug_att_mat / aug_att_mat.sum(dim=-1, keepdim=True)
 
     joint_attentions = torch.zeros_like(aug_att_mat)
     joint_attentions[0] = aug_att_mat[0]
@@ -151,11 +124,11 @@ def predict(image):
     inputs = processor(images=image, return_tensors="pt")
     inputs = {key: value.to(DEVICE) for key, value in inputs.items()}
 
-    # Reset captured attentions before forward pass
     captured_attentions.clear()
 
     with torch.no_grad():
-        outputs = model(**inputs)
+        # Passing output_attentions=True forces ViT layers to return attention matrices
+        outputs = model(**inputs, output_attentions=True)
         logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits
 
     probs = torch.nn.functional.softmax(logits, dim=-1)
@@ -164,30 +137,23 @@ def predict(image):
     prediction = model.config.id2label[predicted_class_idx.item()]
     confidence_pct = round(confidence.item() * 100, 2)
 
-    # Attention rollout using captured attentions
+    # Compute rollout
     rollout = compute_attention_rollout(captured_attentions)
 
     if rollout is not None:
-        mask = rollout[0, 1:]
+        mask = rollout[0, 1:] # Token 0 is CLS token, rest are patch tokens
         size = int(mask.shape[0] ** 0.5)
-        mask = mask.reshape(size, size).cpu().numpy()
+        mask = mask.reshape(size, size).numpy()
 
         mask = cv2.resize(mask, image.size)
         mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
 
-        heatmap = cv2.applyColorMap(
-            np.uint8(255 * mask),
-            cv2.COLORMAP_JET
-        )
-
-        overlay = cv2.addWeighted(
-            np.array(image),
-            0.6,
-            heatmap,
-            0.4,
-            0
-        )
+        heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
+        heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+
+        overlay = cv2.addWeighted(np.array(image), 0.6, heatmap, 0.4, 0)
     else:
+        print("Warning: No attention matrices captured during forward pass.")
         overlay = np.array(image)
 
     return prediction, f"{confidence_pct}%", overlay
@@ -197,68 +163,16 @@ def predict(image):
 # -----------------------------
 
 custom_css = """
-/* Professional Adaptive Theme */
-:root {
-    --primary-blue: #2563eb;
-    --hero-text: #0f172a;
-}
-.dark {
-    --hero-text: #f8fafc;
-}
-body {
-    background-color: var(--background-fill-primary);
-}
-.hero {
-    text-align: center;
-    font-family: 'Inter', sans-serif;
-    font-size: 48px;
-    font-weight: 800;
-    letter-spacing: -0.04em;
-    margin-top: 50px;
-    color: var(--hero-text) !important;
-}
-.sub {
-    text-align: center;
-    opacity: 0.7;
-    font-size: 14px;
-    font-weight: 600;
-    letter-spacing: 0.1em;
-    text-transform: uppercase;
-    margin-bottom: 40px;
-    color: var(--body-text-color);
-}
-.glass {
-    background: var(--block-background-fill) !important;
-    border: 1px solid var(--border-color-primary) !important;
-    border-radius: 12px !important;
-    padding: 24px !important;
-    box-shadow: var(--block-shadow);
-    transition: all 0.2s ease;
-}
-.glass:hover {
-    border-color: var(--primary-blue) !important;
-    box-shadow: 0 4px 20px rgba(37, 99, 235, 0.1);
-}
-button.primary {
-    background: var(--primary-blue) !important;
-    color: white !important;
-    border: none !important;
-    font-weight: 600 !important;
-    padding: 12px 24px !important;
-    border-radius: 8px !important;
-    box-shadow: 0 4px 12px rgba(37, 99, 235, 0.2) !important;
-}
-button.primary:hover {
-    background: #1d4ed8 !important;
-    transform: translateY(-1px);
-    box-shadow: 0 6px 16px rgba(37, 99, 235, 0.3) !important;
-}
-.gr-label {
-    font-weight: 600 !important;
-    font-size: 12px !important;
-    text-transform: uppercase;
-    color: var(--primary-blue) !important;
-}
+:root { --primary-blue: #2563eb; --hero-text: #0f172a; }
+.dark { --hero-text: #f8fafc; }
+body { background-color: var(--background-fill-primary); }
+.hero { text-align: center; font-family: 'Inter', sans-serif; font-size: 48px; font-weight: 800; letter-spacing: -0.04em; margin-top: 50px; color: var(--hero-text) !important; }
+.sub { text-align: center; opacity: 0.7; font-size: 14px; font-weight: 600; letter-spacing: 0.1em; text-transform: uppercase; margin-bottom: 40px; color: var(--body-text-color); }
+.glass { background: var(--block-background-fill) !important; border: 1px solid var(--border-color-primary) !important; border-radius: 12px !important; padding: 24px !important; box-shadow: var(--block-shadow); transition: all 0.2s ease; }
+.glass:hover { border-color: var(--primary-blue) !important; box-shadow: 0 4px 20px rgba(37, 99, 235, 0.1); }
+button.primary { background: var(--primary-blue) !important; color: white !important; border: none !important; font-weight: 600 !important; padding: 12px 24px !important; border-radius: 8px !important; box-shadow: 0 4px 12px rgba(37, 99, 235, 0.2) !important; }
+button.primary:hover { background: #1d4ed8 !important; transform: translateY(-1px); box-shadow: 0 6px 16px rgba(37, 99, 235, 0.3) !important; }
+.gr-label { font-weight: 600 !important; font-size: 12px !important; text-transform: uppercase; color: var(--primary-blue) !important; }
 """
 
 with gr.Blocks(