Deploy auto GPU fallback + FastAPI /predict

.gitattributes

@@ -1,35 +1,35 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,14 +1,14 @@
----
-title: TrueFrame
-emoji: 🚀
-colorFrom: yellow
-colorTo: indigo
-sdk: gradio
-sdk_version: 6.11.0
-app_file: app.py
-pinned: false
-license: apache-2.0
-short_description: A LLaVA-based multimodal classifier with LoRA fine-tuning, d
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: TrueFrame
+emoji: 🚀
+colorFrom: yellow
+colorTo: indigo
+sdk: gradio
+sdk_version: 6.11.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+short_description: A LLaVA-based multimodal classifier with LoRA fine-tuning, d
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,423 +1,464 @@
-import os
-import sys
-import time
-import math
-import json
-from contextlib import nullcontext
-from pathlib import Path
-
-import gradio as gr
-import torch
-from dotenv import load_dotenv
-from PIL import Image, ImageOps
-from transformers import AutoProcessor, AutoImageProcessor, AutoModelForImageClassification
-
-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))
-
-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 = os.getenv("ADAPTER_REPO_ID", "Werrewulf/TMOS-DD")
-ADAPTER_SUBFOLDER = os.getenv("ADAPTER_SUBFOLDER", "")
-
-MODEL_MODE = os.getenv("MODEL_MODE", "cpu-fallback").strip().lower()
-CPU_FALLBACK_MODEL_ID = os.getenv("CPU_FALLBACK_MODEL_ID", "DaMsTaR/Detecto-DeepFake_Image_Detector")
-DEFAULT_INVERT_FALLBACK = CPU_FALLBACK_MODEL_ID.lower() == "damstar/detecto-deepfake_image_detector"
-INVERT_FALLBACK_OUTPUT = os.getenv("INVERT_FALLBACK_OUTPUT", str(DEFAULT_INVERT_FALLBACK)).strip().lower() == "true"
-
-TMOS_PROMPT = "USER: <image>\nIs this video real or produced by AI?\nASSISTANT:"
-TARGET_IMAGE_SIZE = 336
-THRESHOLD = 0.5
-
-model = None
-processor = None
-inference_device = None
-
-
-def resolve_inference_device(model_obj) -> torch.device:
-    if torch.cuda.is_available():
-        return torch.device("cuda")
-
-    device_map = getattr(model_obj, "hf_device_map", None)
-    if isinstance(device_map, dict):
-        for mapped in device_map.values():
-            if isinstance(mapped, str) and mapped.startswith("cuda"):
-                return torch.device(mapped)
-    return torch.device("cpu")
-
-
-def find_classifier_weight_tensor(model_obj):
-    visited = set()
-    queue = [model_obj]
-    while queue:
-        current = queue.pop(0)
-        if current is None:
-            continue
-        obj_id = id(current)
-        if obj_id in visited:
-            continue
-        visited.add(obj_id)
-
-        classifier = getattr(current, "classifier", None)
-        if classifier is not None and hasattr(classifier, "weight"):
-            return classifier.weight
-
-        for attr in ("model", "base_model", "module"):
-            nested = getattr(current, attr, None)
-            if nested is not None:
-                queue.append(nested)
-
-    return None
-
-
-def count_lora_layers(model_obj) -> int:
-    count = 0
-    for _, module in model_obj.named_modules():
-        if hasattr(module, "lora_A") and hasattr(module, "lora_B"):
-            count += 1
-    return count
-
-
-def is_tmos_adapter_config(cfg: dict) -> bool:
-    modules_to_save = cfg.get("modules_to_save") or []
-    target_modules = set(cfg.get("target_modules") or [])
-    required_targets = {"q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"}
-
-    return (
-        "classifier" in modules_to_save
-        and cfg.get("r") == 64
-        and required_targets.issubset(target_modules)
-    )
-
-
-def load_local_adapter_config(adapter_dir: Path) -> dict | None:
-    cfg_path = adapter_dir / "adapter_config.json"
-    if not cfg_path.exists():
-        return None
-    with cfg_path.open("r", encoding="utf-8") as fp:
-        return json.load(fp)
-
-
-def load_remote_adapter_config(repo_id: str, subfolder: str) -> dict | None:
-    from peft import PeftConfig
-
-    try:
-        peft_cfg = PeftConfig.from_pretrained(repo_id, subfolder=subfolder, token=HF_TOKEN)
-        return peft_cfg.to_dict()
-    except Exception:
-        return 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.float32
-
-
-def load_tmos_model():
-    global model, processor, inference_device
-
-    if not torch.cuda.is_available():
-        raise RuntimeError(
-            "TMOS mode requires GPU hardware. Set MODEL_MODE=cpu-fallback for free-tier CPU execution."
-        )
-
-    from peft import PeftModel
-    from tmos_classifier import TMOSClassifier
-
-    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,
-    )
-
-    selected_subfolder = ""
-
-    if local_adapter_file is not None:
-        adapter_source = str(ADAPTER_PATH)
-        local_cfg = load_local_adapter_config(ADAPTER_PATH)
-        if local_cfg is None or not is_tmos_adapter_config(local_cfg):
-            raise RuntimeError(
-                "Local adapter exists but is not TMOS-compatible. Expected modules_to_save=['classifier'], r=64, and TMOS target modules."
-            )
-    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,
-    )
-    base_classifier_weight = find_classifier_weight_tensor(base_model)
-    base_classifier_snapshot = None
-    if base_classifier_weight is not None:
-        base_classifier_snapshot = base_classifier_weight.detach().float().cpu().clone()
-
-    peft_kwargs = {"is_trainable": False, "token": HF_TOKEN}
-    if adapter_source == ADAPTER_REPO_ID:
-        candidate_subfolders = [
-            s for s in [ADAPTER_SUBFOLDER, "multimodal", "multimodal/checkpoint-5", "llava"] if s is not None
-        ]
-
-        last_error = None
-        for subfolder in candidate_subfolders:
-            try:
-                remote_cfg = load_remote_adapter_config(adapter_source, subfolder)
-                if remote_cfg is None or not is_tmos_adapter_config(remote_cfg):
-                    raise ValueError("Adapter config is not TMOS-compatible.")
-
-                current_kwargs = dict(peft_kwargs)
-                if subfolder:
-                    current_kwargs["subfolder"] = subfolder
-                loaded_model = PeftModel.from_pretrained(base_model, adapter_source, **current_kwargs)
-
-                lora_layer_count = count_lora_layers(loaded_model)
-                if lora_layer_count == 0:
-                    raise RuntimeError("Loaded adapter has zero LoRA layers attached.")
-
-                loaded_classifier_weight = find_classifier_weight_tensor(loaded_model)
-                if loaded_classifier_weight is None:
-                    raise RuntimeError("Classifier head not found after adapter load.")
-
-                if base_classifier_snapshot is not None:
-                    classifier_delta = (
-                        loaded_classifier_weight.detach().float().cpu() - base_classifier_snapshot
-                    ).abs().mean().item()
-                    if classifier_delta < 1e-8:
-                        raise RuntimeError(
-                            "Classifier weights did not change after loading adapter; adapter likely incompatible."
-                        )
-
-                model = loaded_model.merge_and_unload()
-                selected_subfolder = subfolder
-                print(
-                    f"Loaded TMOS adapter from repo subfolder: '{subfolder or '.'}' "
-                    f"(lora_layers={lora_layer_count})"
-                )
-                break
-            except Exception as exc:
-                last_error = exc
-                continue
-        else:
-            raise RuntimeError(
-                "No TMOS-compatible adapter found in remote repo. Upload TMOS production weights with classifier head "
-                "(modules_to_save=['classifier'], r=64, 7-target-module LoRA)."
-            ) from last_error
-    else:
-        loaded_model = PeftModel.from_pretrained(base_model, adapter_source, **peft_kwargs)
-        lora_layer_count = count_lora_layers(loaded_model)
-        if lora_layer_count == 0:
-            raise RuntimeError("Local adapter load produced zero LoRA layers attached.")
-        model = loaded_model.merge_and_unload()
-        print(f"Loaded TMOS local adapter (lora_layers={lora_layer_count})")
-
-    model.eval()
-    processor = AutoProcessor.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)
-    processor.patch_size = 14
-    processor.vision_feature_select_strategy = "default"
-    inference_device = resolve_inference_device(model)
-    if adapter_source == ADAPTER_REPO_ID:
-        print(f"TMOS-DD ready on {inference_device} using remote subfolder '{selected_subfolder or '.'}'.")
-    else:
-        print(f"TMOS-DD ready on {inference_device} using local production adapter.")
-
-
-def load_cpu_fallback_model():
-    global model, processor, inference_device
-    print(f"Loading CPU fallback model from {CPU_FALLBACK_MODEL_ID}...")
-
-    processor = AutoImageProcessor.from_pretrained(CPU_FALLBACK_MODEL_ID, token=HF_TOKEN)
-    model = AutoModelForImageClassification.from_pretrained(
-        CPU_FALLBACK_MODEL_ID,
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=True,
-        token=HF_TOKEN,
-    )
-    model.to("cpu").eval()
-    inference_device = torch.device("cpu")
-    print("CPU fallback classifier ready.")
-
-
-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
-
-    if MODEL_MODE == "tmos":
-        load_tmos_model()
-    else:
-        load_cpu_fallback_model()
-
-    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 score_fallback_logits(logits: torch.Tensor, id2label: dict) -> tuple[float, str]:
-    probs = torch.softmax(logits.float(), dim=0)
-
-    fake_indices = []
-    real_indices = []
-    for idx in range(len(probs)):
-        label = str(id2label.get(idx, "")).lower()
-        if any(key in label for key in ["fake", "deepfake", "ai", "synthetic"]):
-            fake_indices.append(idx)
-        if any(key in label for key in ["real", "authentic", "genuine"]):
-            real_indices.append(idx)
-
-    if len(probs) == 2 and not fake_indices and not real_indices:
-        fake_indices = [1]
-        real_indices = [0]
-
-    fake_prob = float(probs[fake_indices].sum().item()) if fake_indices else 0.0
-    real_prob = float(probs[real_indices].sum().item()) if real_indices else 0.0
-
-    total = fake_prob + real_prob
-    if total > 0:
-        prob_fake = fake_prob / total
-    else:
-        prob_fake = float(probs.max().item()) if len(probs) == 1 else float(probs[1].item()) if len(probs) > 1 else 0.5
-
-    if INVERT_FALLBACK_OUTPUT:
-        prob_fake = 1.0 - prob_fake
-
-    label = "Fake" if prob_fake >= THRESHOLD else "Real"
-    return prob_fake, label
-
-
-def infer_image(image: Image.Image):
-    try:
-        if image is None:
-            return None, "Error: please upload an image.", None, None, None, "<div style='color:#f87171;'>Please upload an image before running detection.</div>"
-
-        model_obj, processor_obj, device = load_model_and_processor()
-        prepared_image = preprocess_image(image)
-
-        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:
-            if MODEL_MODE == "tmos":
-                inputs = processor_obj(text=TMOS_PROMPT, images=prepared_image, return_tensors="pt", padding=True)
-                inputs = {name: tensor.to(device) for name, tensor in inputs.items()}
-                outputs = model_obj(
-                    input_ids=inputs["input_ids"],
-                    pixel_values=inputs["pixel_values"],
-                    attention_mask=inputs["attention_mask"],
-                )
-                logit = float(outputs["logit"].squeeze().detach().float().cpu().item())
-                if not math.isfinite(logit):
-                    raise gr.Error("Model produced a non-finite logit (NaN/Inf). Please retry.")
-                prob_fake = float(torch.sigmoid(torch.tensor(logit)).item())
-                label = "Fake" if prob_fake >= THRESHOLD else "Real"
-            else:
-                inputs = processor_obj(images=prepared_image, return_tensors="pt")
-                inputs = {name: tensor.to(device) for name, tensor in inputs.items()}
-                outputs = model_obj(**inputs)
-                logits = outputs.logits.squeeze(0).detach().float().cpu()
-                id2label = getattr(model_obj.config, "id2label", {}) or {}
-                prob_fake, label = score_fallback_logits(logits, id2label)
-
-        if device.type == "cuda":
-            torch.cuda.synchronize()
-
-        elapsed_ms = (time.perf_counter() - start_time) * 1000.0
-        if not math.isfinite(prob_fake):
-            raise gr.Error("Model produced a non-finite probability (NaN/Inf). Please retry.")
-
-        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)
-    except Exception as exc:
-        err = f"Inference failed: {type(exc).__name__}: {exc}"
-        err_html = f"<div style='color:#fca5a5; border:1px solid rgba(252,165,165,0.35); padding:10px; border-radius:10px;'>\n<b>Inference error</b><br>{err}</div>"
-        return None, err, None, None, None, err_html
-
-
-load_model_and_processor()
-
-with gr.Blocks(title="TMOS Deepfake Detector", theme=gr.themes.Soft()) as demo:
-    demo_description = (
-        "TMOS mode (GPU required) enabled.\n\n"
-        if MODEL_MODE == "tmos"
-        else f"CPU fallback mode using {CPU_FALLBACK_MODEL_ID} (invert_output={INVERT_FALLBACK_OUTPUT}).\n\n"
-    )
-    gr.Markdown(
-        "# TMOS Deepfake Detector\n"
-        + demo_description
-        + "> 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")))
+import os
+import sys
+import time
+import math
+import json
+import io
+from contextlib import nullcontext
+from pathlib import Path
+
+import gradio as gr
+import torch
+from fastapi import FastAPI, UploadFile, File
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from dotenv import load_dotenv
+from PIL import Image, ImageOps
+from transformers import AutoProcessor, AutoImageProcessor, AutoModelForImageClassification
+
+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))
+
+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 = os.getenv("ADAPTER_REPO_ID", "Werrewulf/TMOS-DD")
+ADAPTER_SUBFOLDER = os.getenv("ADAPTER_SUBFOLDER", "")
+
+CPU_FALLBACK_MODEL_ID = os.getenv("CPU_FALLBACK_MODEL_ID", "DaMsTaR/Detecto-DeepFake_Image_Detector")
+DEFAULT_INVERT_FALLBACK = CPU_FALLBACK_MODEL_ID.lower() == "damstar/detecto-deepfake_image_detector"
+INVERT_FALLBACK_OUTPUT = os.getenv("INVERT_FALLBACK_OUTPUT", str(DEFAULT_INVERT_FALLBACK)).strip().lower() == "true"
+
+TMOS_PROMPT = "USER: <image>\nIs this video real or produced by AI?\nASSISTANT:"
+TARGET_IMAGE_SIZE = 336
+THRESHOLD = 0.5
+
+model = None
+processor = None
+inference_device = None
+
+
+def resolve_inference_device(model_obj) -> torch.device:
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+
+    device_map = getattr(model_obj, "hf_device_map", None)
+    if isinstance(device_map, dict):
+        for mapped in device_map.values():
+            if isinstance(mapped, str) and mapped.startswith("cuda"):
+                return torch.device(mapped)
+    return torch.device("cpu")
+
+
+def find_classifier_weight_tensor(model_obj):
+    visited = set()
+    queue = [model_obj]
+    while queue:
+        current = queue.pop(0)
+        if current is None:
+            continue
+        obj_id = id(current)
+        if obj_id in visited:
+            continue
+        visited.add(obj_id)
+
+        classifier = getattr(current, "classifier", None)
+        if classifier is not None and hasattr(classifier, "weight"):
+            return classifier.weight
+
+        for attr in ("model", "base_model", "module"):
+            nested = getattr(current, attr, None)
+            if nested is not None:
+                queue.append(nested)
+
+    return None
+
+
+def count_lora_layers(model_obj) -> int:
+    count = 0
+    for _, module in model_obj.named_modules():
+        if hasattr(module, "lora_A") and hasattr(module, "lora_B"):
+            count += 1
+    return count
+
+
+def is_tmos_adapter_config(cfg: dict) -> bool:
+    modules_to_save = cfg.get("modules_to_save") or []
+    target_modules = set(cfg.get("target_modules") or [])
+    required_targets = {"q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"}
+
+    return (
+        "classifier" in modules_to_save
+        and cfg.get("r") == 64
+        and required_targets.issubset(target_modules)
+    )
+
+
+def load_local_adapter_config(adapter_dir: Path) -> dict | None:
+    cfg_path = adapter_dir / "adapter_config.json"
+    if not cfg_path.exists():
+        return None
+    with cfg_path.open("r", encoding="utf-8") as fp:
+        return json.load(fp)
+
+
+def load_remote_adapter_config(repo_id: str, subfolder: str) -> dict | None:
+    from peft import PeftConfig
+
+    try:
+        peft_cfg = PeftConfig.from_pretrained(repo_id, subfolder=subfolder, token=HF_TOKEN)
+        return peft_cfg.to_dict()
+    except Exception:
+        return 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.float32
+
+
+def load_tmos_model():
+    global model, processor, inference_device
+
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "TMOS mode requires GPU hardware. CPU fallback should be used on CPU-only environments."
+        )
+
+    from peft import PeftModel
+    from tmos_classifier import TMOSClassifier
+
+    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,
+    )
+
+    selected_subfolder = ""
+
+    if local_adapter_file is not None:
+        adapter_source = str(ADAPTER_PATH)
+        local_cfg = load_local_adapter_config(ADAPTER_PATH)
+        if local_cfg is None or not is_tmos_adapter_config(local_cfg):
+            raise RuntimeError(
+                "Local adapter exists but is not TMOS-compatible. Expected modules_to_save=['classifier'], r=64, and TMOS target modules."
+            )
+    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,
+    )
+    base_classifier_weight = find_classifier_weight_tensor(base_model)
+    base_classifier_snapshot = None
+    if base_classifier_weight is not None:
+        base_classifier_snapshot = base_classifier_weight.detach().float().cpu().clone()
+
+    peft_kwargs = {"is_trainable": False, "token": HF_TOKEN}
+    if adapter_source == ADAPTER_REPO_ID:
+        candidate_subfolders = [
+            s for s in [ADAPTER_SUBFOLDER, "multimodal", "multimodal/checkpoint-5", "llava"] if s is not None
+        ]
+
+        last_error = None
+        for subfolder in candidate_subfolders:
+            try:
+                remote_cfg = load_remote_adapter_config(adapter_source, subfolder)
+                if remote_cfg is None or not is_tmos_adapter_config(remote_cfg):
+                    raise ValueError("Adapter config is not TMOS-compatible.")
+
+                current_kwargs = dict(peft_kwargs)
+                if subfolder:
+                    current_kwargs["subfolder"] = subfolder
+                loaded_model = PeftModel.from_pretrained(base_model, adapter_source, **current_kwargs)
+
+                lora_layer_count = count_lora_layers(loaded_model)
+                if lora_layer_count == 0:
+                    raise RuntimeError("Loaded adapter has zero LoRA layers attached.")
+
+                loaded_classifier_weight = find_classifier_weight_tensor(loaded_model)
+                if loaded_classifier_weight is None:
+                    raise RuntimeError("Classifier head not found after adapter load.")
+
+                if base_classifier_snapshot is not None:
+                    classifier_delta = (
+                        loaded_classifier_weight.detach().float().cpu() - base_classifier_snapshot
+                    ).abs().mean().item()
+                    if classifier_delta < 1e-8:
+                        raise RuntimeError(
+                            "Classifier weights did not change after loading adapter; adapter likely incompatible."
+                        )
+
+                model = loaded_model.merge_and_unload()
+                selected_subfolder = subfolder
+                print(
+                    f"Loaded TMOS adapter from repo subfolder: '{subfolder or '.'}' "
+                    f"(lora_layers={lora_layer_count})"
+                )
+                break
+            except Exception as exc:
+                last_error = exc
+                continue
+        else:
+            raise RuntimeError(
+                "No TMOS-compatible adapter found in remote repo. Upload TMOS production weights with classifier head "
+                "(modules_to_save=['classifier'], r=64, 7-target-module LoRA)."
+            ) from last_error
+    else:
+        loaded_model = PeftModel.from_pretrained(base_model, adapter_source, **peft_kwargs)
+        lora_layer_count = count_lora_layers(loaded_model)
+        if lora_layer_count == 0:
+            raise RuntimeError("Local adapter load produced zero LoRA layers attached.")
+        model = loaded_model.merge_and_unload()
+        print(f"Loaded TMOS local adapter (lora_layers={lora_layer_count})")
+
+    model.eval()
+    processor = AutoProcessor.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)
+    processor.patch_size = 14
+    processor.vision_feature_select_strategy = "default"
+    inference_device = resolve_inference_device(model)
+    if adapter_source == ADAPTER_REPO_ID:
+        print(f"TMOS-DD ready on {inference_device} using remote subfolder '{selected_subfolder or '.'}'.")
+    else:
+        print(f"TMOS-DD ready on {inference_device} using local production adapter.")
+
+
+def load_cpu_fallback_model():
+    global model, processor, inference_device
+    print(f"Loading CPU fallback model from {CPU_FALLBACK_MODEL_ID}...")
+
+    processor = AutoImageProcessor.from_pretrained(CPU_FALLBACK_MODEL_ID, token=HF_TOKEN)
+    model = AutoModelForImageClassification.from_pretrained(
+        CPU_FALLBACK_MODEL_ID,
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True,
+        token=HF_TOKEN,
+    )
+    model.to("cpu").eval()
+    inference_device = torch.device("cpu")
+    print("CPU fallback classifier ready.")
+
+
+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
+
+    if torch.cuda.is_available():
+        print("GPU detected -> loading TMOS")
+        try:
+            load_tmos_model()
+        except Exception as exc:
+            print(f"TMOS failed: {exc}")
+            print("Falling back to CPU model...")
+            load_cpu_fallback_model()
+    else:
+        print("No GPU detected -> using CPU fallback")
+        load_cpu_fallback_model()
+
+    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 score_fallback_logits(logits: torch.Tensor, id2label: dict) -> tuple[float, str]:
+    probs = torch.softmax(logits.float(), dim=0)
+
+    fake_indices = []
+    real_indices = []
+    for idx in range(len(probs)):
+        label = str(id2label.get(idx, "")).lower()
+        if any(key in label for key in ["fake", "deepfake", "ai", "synthetic"]):
+            fake_indices.append(idx)
+        if any(key in label for key in ["real", "authentic", "genuine"]):
+            real_indices.append(idx)
+
+    if len(probs) == 2 and not fake_indices and not real_indices:
+        fake_indices = [1]
+        real_indices = [0]
+
+    fake_prob = float(probs[fake_indices].sum().item()) if fake_indices else 0.0
+    real_prob = float(probs[real_indices].sum().item()) if real_indices else 0.0
+
+    total = fake_prob + real_prob
+    if total > 0:
+        prob_fake = fake_prob / total
+    else:
+        prob_fake = float(probs.max().item()) if len(probs) == 1 else float(probs[1].item()) if len(probs) > 1 else 0.5
+
+    if INVERT_FALLBACK_OUTPUT:
+        prob_fake = 1.0 - prob_fake
+
+    label = "Fake" if prob_fake >= THRESHOLD else "Real"
+    return prob_fake, label
+
+
+def infer_image(image: Image.Image):
+    try:
+        if image is None:
+            return None, "Error: please upload an image.", None, None, None, "<div style='color:#f87171;'>Please upload an image before running detection.</div>"
+
+        model_obj, processor_obj, device = load_model_and_processor()
+        prepared_image = preprocess_image(image)
+
+        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:
+            if inference_device.type == "cuda":
+                inputs = processor_obj(text=TMOS_PROMPT, images=prepared_image, return_tensors="pt", padding=True)
+                inputs = {name: tensor.to(device) for name, tensor in inputs.items()}
+                outputs = model_obj(
+                    input_ids=inputs["input_ids"],
+                    pixel_values=inputs["pixel_values"],
+                    attention_mask=inputs["attention_mask"],
+                )
+                logit = float(outputs["logit"].squeeze().detach().float().cpu().item())
+                if not math.isfinite(logit):
+                    raise gr.Error("Model produced a non-finite logit (NaN/Inf). Please retry.")
+                prob_fake = float(torch.sigmoid(torch.tensor(logit)).item())
+                label = "Fake" if prob_fake >= THRESHOLD else "Real"
+            else:
+                inputs = processor_obj(images=prepared_image, return_tensors="pt")
+                inputs = {name: tensor.to(device) for name, tensor in inputs.items()}
+                outputs = model_obj(**inputs)
+                logits = outputs.logits.squeeze(0).detach().float().cpu()
+                id2label = getattr(model_obj.config, "id2label", {}) or {}
+                prob_fake, label = score_fallback_logits(logits, id2label)
+
+        if device.type == "cuda":
+            torch.cuda.synchronize()
+
+        elapsed_ms = (time.perf_counter() - start_time) * 1000.0
+        if not math.isfinite(prob_fake):
+            raise gr.Error("Model produced a non-finite probability (NaN/Inf). Please retry.")
+
+        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)
+    except Exception as exc:
+        err = f"Inference failed: {type(exc).__name__}: {exc}"
+        err_html = f"<div style='color:#fca5a5; border:1px solid rgba(252,165,165,0.35); padding:10px; border-radius:10px;'>\n<b>Inference error</b><br>{err}</div>"
+        return None, err, None, None, None, err_html
+
+
+
+api = FastAPI()
+
+api.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@api.post("/predict")
+async def predict(file: UploadFile = File(...)):
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+
+        _, label, prob_fake, confidence, latency, _ = infer_image(image)
+
+        return JSONResponse(
+            {
+                "verdict": label,
+                "confidence_percent": confidence,
+                "p_fake": prob_fake,
+                "latency_ms": latency,
+            }
+        )
+    except Exception as exc:
+        return JSONResponse({"error": str(exc)}, status_code=500)
+
+
+load_model_and_processor()
+
+with gr.Blocks(title="TMOS Deepfake Detector", theme=gr.themes.Soft()) as demo:
+    device_label = "GPU (TMOS Model)" if torch.cuda.is_available() else "CPU Fallback Model"
+    gr.Markdown(
+        f"# TMOS Deepfake Detector\n"
+        f"**Running on:** {device_label}\n\n"
+        f"> 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)
+
+app = gr.mount_gradio_app(api, demo, path="/")
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))

requirements.txt

@@ -1,36 +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==4.36.2
-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
+# 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==4.36.2
+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

@@ -1,216 +1,216 @@
-"""
-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
-
-        # Keep the classifier head in fp32 for numerical stability.
-        self.classifier = nn.Linear(hidden_size, 1, dtype=torch.float32)
-        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, :]
-
-        # Replace non-finite activations defensively before the classifier.
-        pooled = torch.nan_to_num(pooled, nan=0.0, posinf=1e4, neginf=-1e4)
-
-        # Match classifier device to pooled features when model is sharded/offloaded.
-        if self.classifier.weight.device != pooled.device:
-            self.classifier = self.classifier.to(pooled.device)
-
-        # ── 3. Classify ──
-        logit = self.classifier(pooled.float())    # (B, 1)
-        logit = torch.nan_to_num(logit, nan=0.0, posinf=20.0, neginf=-20.0)
-
-        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.")
+"""
+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
+
+        # Keep the classifier head in fp32 for numerical stability.
+        self.classifier = nn.Linear(hidden_size, 1, dtype=torch.float32)
+        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, :]
+
+        # Replace non-finite activations defensively before the classifier.
+        pooled = torch.nan_to_num(pooled, nan=0.0, posinf=1e4, neginf=-1e4)
+
+        # Match classifier device to pooled features when model is sharded/offloaded.
+        if self.classifier.weight.device != pooled.device:
+            self.classifier = self.classifier.to(pooled.device)
+
+        # ── 3. Classify ──
+        logit = self.classifier(pooled.float())    # (B, 1)
+        logit = torch.nan_to_num(logit, nan=0.0, posinf=20.0, neginf=-20.0)
+
+        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.")