Usando from_pretrained_fastai para carga limpia

app.py

@@ -3,7 +3,6 @@ import gradio as gr
 import torch
 import timm
 from PIL import Image
-from torchvision import transforms
 from fastai.vision.all import *
 
 # 1. Categorías (Asegúrate de que el orden sea el mismo que dls.vocab en Colab)
@@ -12,16 +11,15 @@ categories = ['Bug', 'Dark', 'Dragon', 'Electric', 'Fairy', 'Fighting',
               'Normal', 'Poison', 'Psychic', 'Rock', 'Steel', 'Water']
 
 def load_model_fastai_style(weights_path):
-    # En lugar de from_empty, creamos dls mínimos con Datasets vacíos
-    # Esto es universal y no depende de métodos de clase que cambian
+    # Creamos dls mínimos
     empty_ds = Datasets([None], [[]])
     dls = DataLoaders.from_dsets(empty_ds, empty_ds, path='.', bs=1, device='cpu')
     dls.vocab = categories
     
-    # Construimos el Learner
-    learn = vision_learner(dls, 'convnext_tiny', pretrained=False)
+    # Especificamos loss_func explícitamente para evitar el AssertionError
+    learn = vision_learner(dls, 'convnext_tiny', pretrained=False, loss_func=CrossEntropyLossFlat())
     
-    # Cargar pesos con control de errores
+    # Cargar pesos
     state = torch.load(weights_path, map_location='cpu', weights_only=False)
     if isinstance(state, dict) and 'model' in state:
         learn.model.load_state_dict(state['model'])
@@ -31,19 +29,12 @@ def load_model_fastai_style(weights_path):
     learn.model.eval()
     return learn
 
-# Cargamos el modelo
+# Carga
 learn = load_model_fastai_style('checkpoint_1.pth')
 
 def predict(img):
-    # 1. Convertir a PILImage de Fastai
-    img = PILImage.create(img)
-    # 2. Resize manual al tamaño de entrenamiento
-    img = img.resize((126, 126))
-    
-    # 3. Predicción
-    # El método predict gestiona internamente la normalización que el modelo espera
+    img = PILImage.create(img).resize((126, 126))
     pred, pred_idx, probs = learn.predict(img)
-    
     return {categories[i]: float(probs[i]) for i in range(len(categories))}
 
 # Interfaz