import os, sys, types, traceback
import gradio as gr
import torch
from huggingface_hub import snapshot_download

# Importa fastai (con Plum 2) -> fasttransform ya no revienta
from fastai.vision.all import PILImage

# --------- Alias para pickles antiguos que esperan plum.function ---------
try:
    import plum
    from plum import Function  # en Plum 2 existe aquí
    if "plum.function" not in sys.modules:
        m = types.ModuleType("plum.function")
        m.Function = Function
        sys.modules["plum.function"] = m
except Exception:
    print("Error preparando alias plum.function:\n", traceback.format_exc())
    raise

# --------- Loader robusto ---------
def load_fastai_from_hub(repo_id: str, filename: str = "model.pkl"):
    repo_dir = snapshot_download(repo_id)
    pkl_path = os.path.join(repo_dir, filename)
    print("Loading:", pkl_path, "exists:", os.path.exists(pkl_path), "size:", os.path.getsize(pkl_path))

    # PyTorch >=2.6: hay que forzar weights_only=False para cargar objetos completos
    learn = torch.load(pkl_path, map_location="cpu", weights_only=False)
    learn.dls.cpu()
    return learn

MODEL_REPO = "Camayli/practica1"
learner = load_fastai_from_hub(MODEL_REPO, "model.pkl")
labels = learner.dls.vocab

def predict(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = learner.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
)
demo.launch()