import gradio as gr
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch.nn.functional as F

# 1. Загрузка модели
model_id = "najicreator90856/is-it-nsfw_ai-moderator"
print(f"Загрузка модели {model_id}...")

processor = AutoImageProcessor.from_pretrained(model_id)
model = AutoModelForImageClassification.from_pretrained(model_id)
model.eval()

def predict_nsfw(img):
    if img is None:
        return "Пожалуйста, загрузите изображение."
    
    # Подготовка
    inputs = processor(images=img, return_tensors="pt")
    
    # Предсказание
    with torch.no_grad():
        logits = model(**inputs).logits
        probs = F.softmax(logits, dim=-1)[0]
    
    # Формируем словарь для компонента Label
    # Он автоматически покажет полоски вероятности
    results = {
        "SFW": float(probs[0]),
        "NSFW": float(probs[1])
    }
    return results

# 2. Создание интерфейса через Blocks
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🔞 Is It NSFW? AI Moderator")
    gr.Markdown("Быстрая проверка контента на базе ResNet-18 (Точность 99.44%)")
    
    with gr.Row():
        # Левая колонка - Ввод
        with gr.Column(scale=1):
            input_img = gr.Image(
                type="pil", 
                label="Перетащите фото сюда",
                height=400
            )
            submit_btn = gr.Button("Проверить", variant="primary")
            
        # Правая колонка - Результат
        with gr.Column(scale=1):
            output_label = gr.Label(
                num_top_classes=2, 
                label="Результат анализа"
            )
            
            gr.Markdown("""
            ### О модели:
            Данная нейросеть обучена различать NSFW и SFW контент.
            - **SFW**: Безопасно для работы/просмотра.
            - **NSFW**: Контент 18+ (Порнография, Хентай).
            """)

    # Ссылки на донаты
    gr.Markdown("---")
    with gr.Row():
        gr.Markdown("""
        ### Поддержать автора:
        [DonationAlerts](https://www.donationalerts.com/r/najicreator) | 
        **USDT TON:** `UQDcC3o7u4LS07pr-cQI-Gxx3hPjXKrRrNQdGBj0WV7SADe3`
        """)

    # Логика работы кнопки
    submit_btn.click(
        fn=predict_nsfw, 
        inputs=input_img, 
        outputs=output_label
    )

if __name__ == "__main__":
    # Убираем примеры (examples), чтобы не было ошибок с путями файлов
    demo.launch()