import gradio as gr

import torch

import torch.nn as nn

from torchvision import models, transforms

from safetensors.torch import load_file

from PIL import Image



# 1. Rebuild the blank ResNet-50 architecture

model = models.resnet50(weights=None)

model.fc = nn.Linear(model.fc.in_features, 2) # 2 classes: Fake (0) and Real (1)



# 2. Load YOUR trained weights from the safetensors file

# Using map_location='cpu' ensures it works on Hugging Face's free CPU tier!

state_dict = load_file("model.safetensors", device="cpu")



# (Optional safety check) If you trained with DataParallel, keys might have "module." in front.

# This removes it so the weights load perfectly no matter what.

state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}



# Inject the weights into the skeleton and set to test mode

model.load_state_dict(state_dict)

model.eval()



# 3. Define the strict patch transform (No randomness!)

test_transform = transforms.Compose([

    # transforms.Resize(256), # Standard practice to resize slightly before center cropping

    transforms.CenterCrop(224), 

    transforms.ToTensor(),

    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])

])



classes = ['FAKE', 'REAL']



# 4. Create the prediction function that Gradio will call

def predict_image(img):

    if img is None:

        return None

        

    # Ensure image is strictly RGB (drops alpha channels from PNGs)

    img = img.convert('RGB')

    

    # Apply transforms and add the batch dimension

    img_tensor = test_transform(img).unsqueeze(0) 

    

    with torch.no_grad():

        # Pass through model

        preds = model(img_tensor)

        # Convert raw numbers to percentages (0.0 to 1.0)

        probs = torch.nn.functional.softmax(preds[0], dim=0)

        

    # Gradio expects a dictionary of { "Class Name": probability }

    return {classes[i]: float(probs[i]) for i in range(2)}



# 5. Build and launch the Web App!

description_text = """



###  How it works:



This model analyzes patches of an image to determine if an image is real or AI-generated. 







###  Limitations for best results:



* **Resolution Sweet Spot:** Works flawlessly on standard AI resolutions and mid-sized images (from **512x512 up to around 1000x1500 pixels**, like 640x832 or 880x1320). 



* **The 4K Danger Zone:** Ultra-high-resolution (like **3840x2160 / 4K**) images will cause the model to fail. Because the model's 'magnifying glass' is strictly fixed to a 224x224 pixel crop, it ends up looking through a pinhole at less than 0.6% of a 4K image, causing it to lose context and guess randomly.



* **Centered Subjects:** The model strictly scans the dead-center of the image. If the AI artifacts or mistakes (like extra fingers or warped backgrounds) are on the far edges, the model won't see them!



* **No Screenshots:** Heavy compression (like taking a screenshot or downloading from messaging apps) destroys the microscopic forensic evidence. Please upload the raw, original files.



"""







# Inside your interface:



interface = gr.Interface(



    fn=predict_image,



    inputs=gr.Image(type="pil", label="Upload an Image"),



    outputs=gr.Label(num_top_classes=2, label="Prediction"),



    title="FakeOut: AI Image Detector",



    description=description_text, 



    flagging_mode="never"



)



interface.launch()