# Deepfake Explainer Inference Helper
import torch
import os
from PIL import Image
from transformers import AutoProcessor, AutoModelForCausalLM
from peft import PeftModel

# Function to analyze an image for deepfakes
def analyze_deepfake(image_path, custom_prompt=None):
    """Analyze an image for signs of deepfakes with detailed explanation"""
    # Set up model paths
    base_model_id = "unsloth/llama-3.2-11b-vision-instruct"
    adapter_id = "saakshigupta/deepfake-explainer-1"
    
    # Load processor
    processor = AutoProcessor.from_pretrained(base_model_id)
    
    # Load base model
    model = AutoModelForCausalLM.from_pretrained(
        base_model_id,
        torch_dtype=torch.float16,
        device_map="auto"
    )
    
    # Load adapter
    model = PeftModel.from_pretrained(model, adapter_id)
    
    # Load image
    image = Image.open(image_path).convert("RGB")
    
    # Create prompt
    if custom_prompt is None:
        prompt = "Analyze this image carefully and determine if it's a deepfake. Provide both a technical explanation and a simple explanation anyone can understand."
    else:
        prompt = custom_prompt
    
    # Process the image and text
    inputs = processor(text=prompt, images=image, return_tensors="pt")
    
    # Fix cross-attention mask
    if 'cross_attention_mask' in inputs and 0 in inputs['cross_attention_mask'].shape:
        batch_size, seq_len, _, num_tiles = inputs['cross_attention_mask'].shape
        visual_features = 6404  # Critical dimension from training
        new_mask = torch.ones((batch_size, seq_len, visual_features, num_tiles))
        inputs['cross_attention_mask'] = new_mask
        print("Fixed cross-attention mask dimensions")
    
    # Move to device
    inputs = {k: v.to(model.device) for k, v in inputs.items() if isinstance(v, torch.Tensor)}
    
    # Generate output
    print("Generating analysis...")
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            top_p=0.9
        )
    
    # Decode output
    response = processor.decode(output_ids[0], skip_special_tokens=True)
    
    # Extract the model's response (removing the prompt)
    if prompt in response:
        result = response.split(prompt)[-1].strip()
    else:
        result = response
    
    return result

# Example usage
if __name__ == "__main__":
    # Example image path
    image_path = input("Enter path to image for deepfake analysis: ")
    
    # Analyze image
    analysis = analyze_deepfake(image_path)
    
    # Print result
    print("
===== DEEPFAKE ANALYSIS RESULT =====
")
    print(analysis)