import torch
import torchvision.transforms as transforms
from PIL import Image
import gradio as gr
from timm import create_model
import torch.nn as nn
import os

# ==== Your custom model ====
class VisionTransformer(nn.Module):
    def __init__(self, num_classes, model_name):
        super(VisionTransformer, self).__init__()
        self.model = create_model(model_name, pretrained=False, num_classes=num_classes)
        self.model.head = nn.Sequential(
            nn.Linear(self.model.num_features, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        return self.model(x)

# ==== Load the trained weights ====
model_path = "./models/vit_small_patch16_224_best.pth"
device = torch.device("cpu")  # MPS not available on HF Spaces

model = VisionTransformer(num_classes=2, model_name="vit_small_patch16_224")
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()

# ==== Image transforms ====
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
])

def predict(img: Image.Image):
    img_tensor = transform(img).unsqueeze(0)
    with torch.no_grad():
        outputs = model(img_tensor)
        _, pred = torch.max(outputs, 1)
        label = "Fake" if pred.item() == 0 else "Real"
        return label

gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="label",
    title="Luxury Item Checker",
    description="Upload an image to check if it's a Real or Fake item."
).launch()