import gradio as gr
import torch
from PIL import Image
from torchvision import transforms
from transformers import AutoConfig, AutoModelForImageClassification
import sys
from huggingface_hub import snapshot_download

# Download model
model_path = snapshot_download("shahad-alh/arabichar-finetuned-v2")
sys.path.append(model_path)

# Load model
config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForImageClassification.from_pretrained(model_path, config=config, trust_remote_code=True)
model.eval()

# Preprocessing
transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),
    transforms.Resize((32, 32)),
    transforms.ToTensor()
])

# Prediction function
def predict(image: Image.Image):
    try:
        print("🔍 Received image:", image)

        tensor = transform(image).unsqueeze(0)  # Add batch dimension
        print("📦 Transformed tensor shape:", tensor.shape)

        with torch.no_grad():
            logits = model(tensor).logits
            predicted = torch.argmax(logits, dim=1).item()
            print("✅ Prediction index:", predicted)

            label = config.id2label[str(predicted)]
            print("🏷️ Label:", label)
        return label

    except Exception as e:
        print("❌ Prediction error:", e)
        return "خطأ"


# ✅ Prediction interface
gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(),
    flagging_mode="never"
).queue().launch(share=True)