Usando from_pretrained_fastai para carga limpia

app.py

@@ -1,49 +1,64 @@
 
 import gradio as gr
-from fastai.vision.all import *
+import gradio as gr
+import torch
 import timm
+from PIL import Image
+from torchvision import transforms
 
-# 1. Lista de categorías exacta
+# 1. Lista de categorías (Asegúrate de que el orden sea el mismo que en tu vocab de Colab)
 categories = ['Bug', 'Dark', 'Dragon', 'Electric', 'Fairy', 'Fighting', 
               'Fire', 'Flying', 'Ghost', 'Grass', 'Ground', 'Ice', 
               'Normal', 'Poison', 'Psychic', 'Rock', 'Steel', 'Water']
 
-def load_clean_model(weights_path):
-    # Creamos la arquitectura ConvNeXt Tiny
+# 2. Carga del modelo directamente en PyTorch
+def load_model(weights_path):
+    # Creamos la arquitectura exacta
     model = timm.create_model('convnext_tiny', pretrained=False, num_classes=len(categories))
     
-    # FORMA COMPATIBLE: Creamos DataLoaders vacíos manualmente
-    # Esto define que el vocabulario son nuestras categorías
-    dls = DataLoaders(Datasets([], [[]]), vocab=categories)
-    
-    # Creamos el Learner
-    learn = Learner(dls, model, loss_func=CrossEntropyLossFlat())
+    # Cargamos los pesos del archivo .pth
+    # weights_only=False es necesario para cargar el estado guardado por fastai
+    state_dict = torch.load(weights_path, map_location='cpu', weights_only=False)
     
-    # Cargamos los pesos (.pth)
-    # Fastai busca el archivo weights_path + '.pth'
-    learn.load(weights_path) 
-    return learn
+    # Si los pesos vienen de un Learner de fastai, suelen estar dentro de una llave 'model'
+    if 'model' in state_dict:
+        state_dict = state_dict['model']
+        
+    # Cargamos los pesos en la arquitectura
+    model.load_state_dict(state_dict, strict=False)
+    model.eval() # Modo evaluación
+    return model
 
-# 2. Inicializar el modelo
-learn = load_clean_model('checkpoint_1')
+# Cargamos el modelo (el archivo debe estar en el repo como checkpoint_1.pth)
+model = load_model('checkpoint_1.pth')
 
+# 3. Función de predicción manual
 def predict(img):
-    # Preparar la imagen
-    img = PILImage.create(img).resize((126, 126))
+    # Preprocesamiento: Resize(126) + Convertir a Tensor + Normalización ImageNet
+    # Esto es exactamente lo que hace fastai por debajo
+    transform = transforms.Compose([
+        transforms.Resize((126, 126)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+    
+    # Convertimos la imagen de Gradio a Tensor
+    img_tensor = transform(img).unsqueeze(0) # Añadir dimensión de batch
     
-    # En inferencia pura (sin dblock), usamos el modelo directamente o learn.predict
-    # learn.predict aplicará el preprocesamiento necesario
-    pred, pred_idx, probs = learn.predict(img)
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        probs = torch.nn.functional.softmax(outputs[0], dim=0)
     
+    # Retornamos diccionario con las probabilidades
     return {categories[i]: float(probs[i]) for i in range(len(categories))}
 
-# 3. Interfaz de Gradio
+# 4. Interfaz de Gradio
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(num_top_classes=3),
     title="Pokemon Type Classifier",
-    description="Inferencia robusta con Fastai y ConvNeXt Tiny."
+    description="Inferencia directa usando PyTorch y ConvNeXt Tiny."
 )
 
 demo.launch()