app.py

import os
import io
import time
import asyncio
import numpy as np
import torch
import torch.nn.functional as F
import cv2
from PIL import Image
from fastapi import FastAPI, UploadFile, File, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from huggingface_hub import snapshot_download, login

from transformers import (
    BlipProcessor, BlipForConditionalGeneration,
    ViTImageProcessor, AutoProcessor, AutoModelForCausalLM,
    CLIPModel, CLIPProcessor, BitsAndBytesConfig
)

app = FastAPI(title="XAI Auditor: Pure Greedy Fast Ensemble")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["X-Processing-Time"]
)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MODELS = {}

@app.on_event("startup")
async def startup_event():
    global MODELS
    token = os.getenv("HF_Token")
    if token: login(token=token)
    
    print("Pinning 8-bit quantized models to memory space...")
    local_dir = snapshot_download(repo_id="SaniaE/Image_Captioning_Ensemble", token=token, local_dir="weights")

    quantization_config = BitsAndBytesConfig(
        load_in_8bit=True,
        llm_int8_threshold=6.0
    ) if DEVICE == "cuda" else None

    # 1. Load Compressed BLIP-Large
    MODELS["blip"] = {
        "model": BlipForConditionalGeneration.from_pretrained(
            os.path.join(local_dir, "blip"),
            quantization_config=quantization_config,
            device_map="auto" if DEVICE == "cuda" else None
        ),
        "processor": BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
    }

    # 2. Load Compressed ViT Track
    MODELS["vit"] = {
        "model": AutoModelForCausalLM.from_pretrained(
            os.path.join(local_dir, "vit"),
            quantization_config=quantization_config,
            device_map="auto" if DEVICE == "cuda" else None
        ),
        "processor": (
            ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning"),
            AutoProcessor.from_pretrained("microsoft/git-large")
        )
    }

    # 3. Load Pinned CLIP Jury
    clip_dtype = torch.float16 if DEVICE == "cuda" else torch.float32
    clip_model = CLIPModel.from_pretrained(os.path.join(local_dir, "clip/clip_model"))
    MODELS["clip"] = {
        "model": clip_model.to(device=DEVICE, dtype=clip_dtype),
        "processor": CLIPProcessor.from_pretrained(os.path.join(local_dir, "clip/clip_processor"))
    }
    
    print("All system weights safely pinned. Pure greedy acceleration paths ready.")

# --- Ultra-Fast Pure Greedy Generation Engine ---

def _generate_balanced_4_track(image, max_len=15):
    """
    Generates exactly 4 captions using clean, single-pass greedy tracks 
    with a repetition penalty to guarantee speed and output variation.
    """
    captions = []
    
    with torch.inference_mode():
        # Track A: Fast Greedy BLIP Pass (2 parallel unique streams)
        b_data = MODELS["blip"]
        b_inputs = b_data["processor"](images=image, return_tensors="pt")
        b_pixels = b_inputs.pixel_values.to(DEVICE)
        batched_b_pixels = b_pixels.repeat(2, 1, 1, 1)
        
        b_ids = b_data["model"].generate(
            pixel_values=batched_b_pixels,
            max_new_tokens=max_len,
            do_sample=False,              # Pure deterministic greedy path
            num_beams=1,                  # Completely eliminate branching tree overhead
            repetition_penalty=1.2,       # Forces variation across the two streams
            early_stopping=True,
            use_cache=True
        )
        b_caps = b_data["processor"].batch_decode(b_ids, skip_special_tokens=True)
        captions.extend([cap.strip() for cap in b_caps])

        # Track B: Fast Greedy ViT Pass (2 parallel unique streams)
        v_data = MODELS["vit"]
        i_proc, t_proc = v_data["processor"]
        v_inputs = i_proc(images=image, return_tensors="pt")
        v_pixels = v_inputs.pixel_values.to(DEVICE)
        batched_v_pixels = v_pixels.repeat(2, 1, 1, 1)
        
        if hasattr(v_inputs, "attention_mask") and v_inputs.attention_mask is not None:
            batched_mask = v_inputs.attention_mask.to(DEVICE).repeat(2, 1)
        else:
            batched_mask = None

        v_ids = v_data["model"].generate(
            pixel_values=batched_v_pixels,
            attention_mask=batched_mask,
            max_new_tokens=max_len,
            do_sample=False,              # Pure deterministic greedy path
            num_beams=1,                  # Completely eliminate branching tree overhead
            repetition_penalty=1.2,       # Forces variation across the two streams
            early_stopping=True,
            use_cache=True
        )
        v_caps = t_proc.batch_decode(v_ids, skip_special_tokens=True)
        captions.extend([cap.strip() for cap in v_caps])
            
    return captions

# --- Endpoints ---

@app.post("/generate")
async def generate_captions(file: UploadFile = File(...)):
    """Generates 4 diverse captions split evenly across architectures for UI balance."""
    start_time = time.perf_counter()
    image = Image.open(file.file).convert("RGB")
    
    # Run the accelerated greedy pipeline
    captions = await asyncio.to_thread(_generate_balanced_4_track, image, 15)
    
    elapsed_time = time.perf_counter() - start_time
    print(f"[BENCHMARK] /generate 4-caption turnaround: {elapsed_time:.4f}s")
    
    return {
        "captions": captions, 
        "metadata": {
            "models_used": ["blip", "blip", "vit", "vit"], 
            "processing_time_sec": round(elapsed_time, 4)
        }
    }

@app.post("/saliency")
async def get_vision_saliency(file: UploadFile = File(...)):
    """Objective Saliency: Native vision encoder self-attention mapping matrix."""
    start_time = time.perf_counter()
    image_bytes = await file.read()
    orig_img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    
    blip = MODELS["blip"]
    inputs = blip["processor"](images=orig_img, return_tensors="pt")
    pixel_values = inputs.pixel_values.to(DEVICE)
    
    with torch.inference_mode():
        outputs = blip["model"].vision_model(pixel_values, output_attentions=True)
        attentions = outputs.attentions[-1] 
        mask_1d = attentions[0, :, 0, 1:].mean(dim=0)
        grid_size = int(np.sqrt(mask_1d.shape[-1]))
        mask = mask_1d.view(grid_size, grid_size).cpu().numpy()

    mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
    w, h = orig_img.size
    mask_resized = cv2.resize(mask, (w, h), interpolation=cv2.INTER_CUBIC)
    mask_blurred = cv2.GaussianBlur(mask_resized, (21, 21), 0)
    
    heatmap_uint8 = np.uint8(255 * mask_blurred)
    heatmap_bgr = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_MAGMA)
    heatmap_rgb = cv2.cvtColor(heatmap_bgr, cv2.COLOR_BGR2RGB)
    
    blended_np = cv2.addWeighted(np.array(orig_img), 0.5, heatmap_rgb, 0.5, 0)
    blended_img = Image.fromarray(blended_np)
    
    buf = io.BytesIO()
    blended_img.save(buf, format="PNG")
    buf.seek(0)
    
    return StreamingResponse(buf, media_type="image/png")

@app.post("/audit")
async def internal_debate_audit(file: UploadFile = File(...), user_prompt: str = Query(...)):
    """The CLIP-Powered Jury: Decoupled visual alignment auditing pass."""
    start_time = time.perf_counter()
    image_bytes = await file.read()
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    
    blip_caption = (await asyncio.to_thread(_generate_balanced_4_track, image, 15))[0]
    
    clip_m = MODELS["clip"]["model"]
    clip_p = MODELS["clip"]["processor"]
    clip_dtype = torch.float16 if DEVICE == "cuda" else torch.float32
    
    image_inputs = clip_p(images=image, return_tensors="pt")
    text_inputs = clip_p(text=[user_prompt, blip_caption], return_tensors="pt", padding=True)
    
    with torch.inference_mode():
        img_pixels = image_inputs.pixel_values.to(device=DEVICE, dtype=clip_dtype)
        txt_ids = text_inputs.input_ids.to(DEVICE)
        txt_mask = text_inputs.attention_mask.to(DEVICE)
        
        image_features = clip_m.get_image_features(pixel_values=img_pixels)
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        
        text_features = clip_m.get_text_features(input_ids=txt_ids, attention_mask=txt_mask)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
        
        logits_per_image = (image_features @ text_features.T) * clip_m.logit_scale.exp()
        probs = F.softmax(logits_per_image, dim=-1).cpu().to(torch.float32).numpy()[0]
    
    u_score, m_score = float(probs[0]), float(probs[1])
    verdict = "Model Bias Detected." if abs(u_score - m_score) >= 0.15 else "Consensus: High Alignment."
    if u_score < 0.35: verdict = "Perspective Divergence: Intent not grounded in image."

    return {
        "perspectives": {"user": user_prompt, "ai": blip_caption},
        "audit_scores": {"intent_grounding": round(u_score, 4), "ai_grounding": round(m_score, 4)},
        "verdict": verdict,
        "metadata": {"processing_time_sec": round(time.perf_counter() - start_time, 4)}
    }