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)} }