import gradio as gr
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image

# Load model + processor (auto cached inside Spaces)
processor = AutoImageProcessor.from_pretrained("prithivMLmods/Realistic-Gender-Classification")
model = AutoModelForImageClassification.from_pretrained("prithivMLmods/Realistic-Gender-Classification")

def predict(image):
    # Preprocess
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0].cpu().numpy()

    # Labels
    labels = list(model.config.id2label.values())

    # Clean dict for FlutterFlow
    result = {
        "female": float(probs[labels.index("female portrait")]),
        "male": float(probs[labels.index("male portrait")])
    }
    return result

# Gradio interface (Spaces auto-hosts this)
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON()
)

if __name__ == "__main__":
    demo.launch()