Add TMOS deepfake Gradio Space

app.py

@@ -0,0 +1,187 @@
+import os
+import sys
+import time
+from contextlib import nullcontext
+from pathlib import Path
+
+import gradio as gr
+import torch
+from dotenv import load_dotenv
+from peft import PeftModel
+from PIL import Image, ImageOps
+from transformers import AutoProcessor
+
+ROOT_DIR = Path(__file__).resolve().parent
+SCRIPTS_DIR = ROOT_DIR / "scripts"
+
+if str(SCRIPTS_DIR) not in sys.path:
+    sys.path.insert(0, str(SCRIPTS_DIR))
+
+from tmos_classifier import TMOSClassifier
+
+load_dotenv()
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+BASE_MODEL_ID = "llava-hf/llava-1.5-7b-hf"
+ADAPTER_PATH = ROOT_DIR / "final-production-weights" / "best_model"
+ADAPTER_REPO_ID = "Werrewulf/tmos-dd-weights"
+PROMPT = "USER: <image>\nIs this image real or produced by AI?\nASSISTANT:"
+TARGET_IMAGE_SIZE = 336
+THRESHOLD = 0.5
+
+model = None
+processor = None
+inference_device = None
+
+
+def select_torch_dtype() -> torch.dtype:
+    if torch.cuda.is_available():
+        return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    return torch.float16
+
+
+def load_model_and_processor():
+    global model, processor, inference_device
+
+    if model is not None and processor is not None and inference_device is not None:
+        return model, processor, inference_device
+
+    adapter_source = None
+    local_adapter_file = next(
+        (
+            candidate
+            for candidate in (
+                ADAPTER_PATH / "adapter_model.safetensors",
+                ADAPTER_PATH / "adapter_model.bin",
+            )
+            if candidate.exists()
+        ),
+        None,
+    )
+
+    if local_adapter_file is not None:
+        adapter_source = str(ADAPTER_PATH)
+    else:
+        adapter_source = ADAPTER_REPO_ID
+
+    dtype = select_torch_dtype()
+    print(f"Loading TMOS-DD model from {adapter_source} with dtype={dtype}...")
+
+    base_model = TMOSClassifier(
+        base_model_id=BASE_MODEL_ID,
+        torch_dtype=dtype,
+        device_map="auto",
+        token=HF_TOKEN,
+    )
+    model = PeftModel.from_pretrained(base_model, adapter_source, is_trainable=False, token=HF_TOKEN)
+    model.eval()
+
+    processor = AutoProcessor.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)
+    processor.patch_size = 14
+    processor.vision_feature_select_strategy = "default"
+
+    inference_device = next(model.parameters()).device
+    print(f"TMOS-DD ready on {inference_device}.")
+    return model, processor, inference_device
+
+
+def preprocess_image(image: Image.Image) -> Image.Image:
+    image = image.convert("RGB")
+    return ImageOps.contain(image, (TARGET_IMAGE_SIZE, TARGET_IMAGE_SIZE), method=Image.Resampling.BICUBIC)
+
+
+def confidence_card(prob_fake: float, label: str) -> str:
+    confidence = prob_fake if label == "Fake" else 1.0 - prob_fake
+    confidence_pct = confidence * 100.0
+    fake_pct = prob_fake * 100.0
+    real_pct = (1.0 - prob_fake) * 100.0
+    accent = "#ef4444" if label == "Fake" else "#10b981"
+
+    return f"""
+    <div style="border:1px solid rgba(255,255,255,0.12); border-radius:16px; padding:16px; background:linear-gradient(135deg, rgba(17,24,39,0.92), rgba(15,23,42,0.96)); color:white;">
+      <div style="font-size:0.85rem; opacity:0.8; letter-spacing:0.04em; text-transform:uppercase; margin-bottom:8px;">Confidence</div>
+      <div style="display:flex; align-items:baseline; gap:10px; margin-bottom:12px;">
+        <div style="font-size:2rem; font-weight:700; color:{accent};">{confidence_pct:.2f}%</div>
+        <div style="font-size:1rem; opacity:0.9;">for <strong>{label}</strong></div>
+      </div>
+      <div style="height:12px; width:100%; background:rgba(255,255,255,0.08); border-radius:999px; overflow:hidden; margin-bottom:10px;">
+        <div style="height:100%; width:{fake_pct:.2f}%; background:linear-gradient(90deg, #f87171, #ef4444);"></div>
+      </div>
+      <div style="display:flex; justify-content:space-between; font-size:0.9rem; opacity:0.95;">
+        <span>Real: {real_pct:.2f}%</span>
+        <span>Fake: {fake_pct:.2f}%</span>
+      </div>
+    </div>
+    """
+
+
+def infer_image(image: Image.Image):
+    if image is None:
+        raise gr.Error("Please upload an image before running detection.")
+
+    model_obj, processor_obj, device = load_model_and_processor()
+
+    prepared_image = preprocess_image(image)
+    inputs = processor_obj(text=PROMPT, images=prepared_image, return_tensors="pt", padding=True)
+    inputs = {name: tensor.to(device) for name, tensor in inputs.items()}
+
+    autocast_context = (
+        torch.autocast(device_type="cuda", dtype=select_torch_dtype())
+        if device.type == "cuda"
+        else nullcontext()
+    )
+
+    start_time = time.perf_counter()
+    with torch.inference_mode(), autocast_context:
+        outputs = model_obj(
+            input_ids=inputs["input_ids"],
+            pixel_values=inputs["pixel_values"],
+            attention_mask=inputs["attention_mask"],
+        )
+
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+
+    elapsed_ms = (time.perf_counter() - start_time) * 1000.0
+
+    logit = float(outputs["logit"].squeeze().detach().float().cpu().item())
+    prob_fake = float(torch.sigmoid(torch.tensor(logit)).item())
+    label = "Fake" if prob_fake >= THRESHOLD else "Real"
+    confidence = prob_fake if label == "Fake" else 1.0 - prob_fake
+
+    return prepared_image, label, round(prob_fake, 6), round(confidence * 100.0, 2), round(elapsed_ms, 2), confidence_card(prob_fake, label)
+
+
+load_model_and_processor()
+
+with gr.Blocks(title="TMOS Deepfake Detector", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        "# TMOS Deepfake Detector\n"
+        "Research demo for image-based deepfake detection with a deterministic classification head on top of LLaVA-1.5-7B.\n\n"
+        "> Warning: runs on free infrastructure, so startup and inference may take time."
+    )
+
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="Upload image")
+        with gr.Column():
+            prediction_output = gr.Textbox(label="Prediction", interactive=False)
+            probability_output = gr.Number(label="P(fake)", interactive=False, precision=6)
+            confidence_output = gr.Number(label="Confidence (%)", interactive=False, precision=2)
+            latency_output = gr.Number(label="Latency (ms)", interactive=False, precision=2)
+
+    preview_output = gr.Image(label="Processed image passed to the model", interactive=False)
+    confidence_html = gr.HTML()
+
+    detect_button = gr.Button("Run detection", variant="primary")
+
+    detect_button.click(
+        fn=infer_image,
+        inputs=image_input,
+        outputs=[preview_output, prediction_output, probability_output, confidence_output, latency_output, confidence_html],
+    )
+
+demo.queue(default_concurrency_limit=1, max_size=8)
+
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))

requirements.txt

@@ -0,0 +1,36 @@
+# Core ML
+numpy>=1.24.0
+python-dotenv
+gradio
+torch>=2.0.0
+torchvision>=0.15.0
+torchaudio>=2.0.0
+torchcodec
+Pillow
+
+# Dependencies
+albumentations>=0.5.2
+datasets
+huggingface_hub
+scikit-learn>=1.3.0
+scikit-image>=0.21.0
+pandas>=2.0.0
+matplotlib>=3.7.0
+seaborn
+transformers
+peft
+accelerate
+diffusers
+opencv-python
+
+# Optional fallback for lower-memory GPU execution
+bitsandbytes
+
+# M2TR specific
+yacs==0.1.8
+nbconvert
+tensorboard==2.20.0
+tqdm==4.67.1
+PyYAML==6.0.3
+simplejson==3.20.2
+fvcore

scripts/tmos_classifier.py

@@ -0,0 +1,208 @@
+"""
+TMOS_Classifier: Binary classification head on top of LLaVA's transformer backbone.
+
+Strips the autoregressive lm_head and replaces it with a single nn.Linear(hidden_size, 1)
+for binary deepfake detection (0 = Real, 1 = Fake).
+
+Usage:
+    from tmos_classifier import TMOSClassifier, TMOS_LORA_CONFIG
+    
+    classifier = TMOSClassifier(base_model_id="llava-hf/llava-1.5-7b-hf")
+    classifier = get_peft_model(classifier, TMOS_LORA_CONFIG)
+    
+    logit = classifier(input_ids=..., pixel_values=..., attention_mask=...)
+    loss  = nn.BCEWithLogitsLoss()(logit, label)
+"""
+
+import torch
+import torch.nn as nn
+from transformers import LlavaForConditionalGeneration
+from peft import LoraConfig
+
+
+# ─── LoRA Configuration ──────────────────────────────────────────────
+# Massive expansion: r=64 across ALL linear layers in the LLM backbone.
+# We exclude lm_head (we discard it), fc1/fc2/out_proj (CLIP vision),
+# and linear_1/linear_2 (multi-modal projector) from LoRA to keep
+# the vision encoder frozen and only adapt the language transformer.
+
+TMOS_LORA_CONFIG = LoraConfig(
+    r=64,
+    lora_alpha=128,               # 2x rank as a common heuristic
+    target_modules=[
+        "q_proj", "k_proj", "v_proj", "o_proj",
+        "gate_proj", "up_proj", "down_proj",
+    ],
+    lora_dropout=0.1,
+    bias="none",
+    task_type=None,                # Custom classifier — not a causal LM
+    modules_to_save=["classifier"],  # Always train the classification head
+)
+
+
+class TMOSClassifier(nn.Module):
+    """
+    Binary classifier built on the LLaVA transformer backbone.
+    
+    Architecture:
+        pixel_values ──► CLIP Vision Tower ──► Multi-Modal Projector ──┐
+                                                                       ├──► LLaMA Transformer ──► last_hidden_state[:, -1, :] ──► classifier ──► logit
+        input_ids ──► Token Embedding ─────────────────────────────────┘
+    
+    The lm_head is never used. We extract the final token's hidden state 
+    and pass it through a learned nn.Linear(hidden_size, 1) head.
+    """
+
+    def __init__(self, base_model_id, torch_dtype=torch.float16, device_map="auto", token=None):
+        super().__init__()
+
+        # Load the full LLaVA model (we need vision tower + projector + LLM)
+        self.base = LlavaForConditionalGeneration.from_pretrained(
+            base_model_id,
+            torch_dtype=torch_dtype,
+            low_cpu_mem_usage=True,
+            device_map=device_map,
+            token=token,
+        )
+
+        hidden_size = self.base.config.text_config.hidden_size  # 4096 for 7B
+
+        # Freeze the lm_head — we won't use it, but freezing prevents
+        # wasted gradient computation if PEFT accidentally wraps it.
+        for param in self.base.lm_head.parameters():
+            param.requires_grad = False
+
+        # Binary classification head: hidden_size → 1 logit
+        self.classifier = nn.Linear(hidden_size, 1, dtype=torch_dtype)
+        nn.init.xavier_uniform_(self.classifier.weight)
+        nn.init.zeros_(self.classifier.bias)
+
+    def forward(
+        self,
+        input_ids=None,
+        pixel_values=None,
+        attention_mask=None,
+        labels=None,            # float tensor of shape (B,) — 0.0=real, 1.0=fake
+        **kwargs,               # absorb extra keys from data collator
+    ):
+        """
+        Single deterministic forward pass → logit + optional BCE loss.
+        
+        Returns:
+            dict with keys:
+                "logit":  (B, 1) raw logit  
+                "loss":   scalar BCE loss (only if labels provided)
+        """
+        # ── 1. Forward through the LLaVA backbone ──
+        # We call the internal model (vision + projector + LLM) directly,
+        # asking for hidden states, NOT for language-model logits.
+        outputs = self.base.model(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            attention_mask=attention_mask,
+            return_dict=True,
+        )
+
+        # last_hidden_state: (B, seq_len, hidden_size)
+        last_hidden_state = outputs.last_hidden_state
+
+        # ── 2. Pool: extract the final non-padded token per sequence ──
+        if attention_mask is not None:
+            # Sum of mask gives the sequence length (excluding padding)
+            # Index of the last real token = seq_lengths - 1
+            seq_lengths = attention_mask.sum(dim=1).long() - 1
+            # Clamp to valid range
+            seq_lengths = seq_lengths.clamp(min=0, max=last_hidden_state.size(1) - 1)
+            # Gather the hidden state at each sequence's last real token
+            pooled = last_hidden_state[
+                torch.arange(last_hidden_state.size(0), device=last_hidden_state.device),
+                seq_lengths,
+            ]
+        else:
+            # No mask → just take the last position
+            pooled = last_hidden_state[:, -1, :]
+
+        # ── 3. Classify ──
+        logit = self.classifier(pooled)    # (B, 1)
+
+        result = {"logit": logit}
+
+        # ── 4. Loss ──
+        if labels is not None:
+            labels = labels.to(logit.dtype).to(logit.device)
+            if labels.dim() == 1:
+                labels = labels.unsqueeze(1)   # (B,) → (B, 1)
+            loss_fn = nn.BCEWithLogitsLoss()
+            result["loss"] = loss_fn(logit, labels)
+
+        return result
+
+    def prepare_inputs_for_generation(self, *args, **kwargs):
+        """Stub required by PEFT — we never generate text."""
+        raise NotImplementedError("TMOSClassifier does not support generation.")
+
+    def gradient_checkpointing_enable(self, **kwargs):
+        """Delegate to the base model for HF Trainer compatibility."""
+        self.base.model.gradient_checkpointing_enable(**kwargs)
+
+    @property
+    def config(self):
+        """Expose the base model config for PEFT."""
+        return self.base.config
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property 
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+
+# ─── Standalone Test ──────────────────────────────────────────────────
+if __name__ == "__main__":
+    import os
+    from dotenv import load_dotenv
+    load_dotenv()
+    HF_TOKEN = os.getenv("HF_TOKEN")
+
+    print("Testing TMOSClassifier...")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    clf = TMOSClassifier(
+        base_model_id="llava-hf/llava-1.5-7b-hf",
+        torch_dtype=torch.float16,
+        token=HF_TOKEN,
+    )
+    clf.to(device)
+
+    # Print parameter counts
+    total = sum(p.numel() for p in clf.parameters())
+    trainable = sum(p.numel() for p in clf.parameters() if p.requires_grad)
+    print(f"Total params:     {total:>12,}")
+    print(f"Trainable params: {trainable:>12,}")
+    print(f"Classifier head:  {sum(p.numel() for p in clf.classifier.parameters()):,}")
+
+    # Smoke test with dummy input
+    from transformers import AutoProcessor
+    processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf", token=HF_TOKEN)
+    processor.patch_size = 14
+    processor.vision_feature_select_strategy = "default"
+
+    from PIL import Image
+    dummy_img = Image.new("RGB", (336, 336), color=(128, 128, 128))
+    inputs = processor(
+        text="USER: <image>\nIs this real?\nASSISTANT:",
+        images=dummy_img,
+        return_tensors="pt",
+    ).to(device)
+
+    labels = torch.tensor([1.0], device=device)  # fake
+
+    with torch.no_grad():
+        out = clf(**inputs, labels=labels)
+    
+    print(f"Logit:  {out['logit'].item():.4f}")
+    print(f"Loss:   {out['loss'].item():.4f}")
+    print(f"Prob:   {torch.sigmoid(out['logit']).item():.4f}")
+    print("Test passed.")