Usando from_pretrained_fastai para carga limpia

app.py

@@ -9,16 +9,18 @@ categories = ['Bug', 'Dark', 'Dragon', 'Electric', 'Fairy', 'Fighting',
               'Normal', 'Poison', 'Psychic', 'Rock', 'Steel', 'Water']
 
 def load_clean_model(weights_path):
-    # Reconstruimos la arquitectura ConvNeXt Tiny
-    # num_classes debe coincidir con el número de categorías
+    # Creamos la arquitectura ConvNeXt Tiny
     model = timm.create_model('convnext_tiny', pretrained=False, num_classes=len(categories))
     
-    # Creamos un Learner mínimo. Al usar dls.from_empty no hay procesos de carga de datos.
-    dls = DataLoaders.from_empty(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 (.pth)
-    # Fastai buscará el archivo weights_path + '.pth' (checkpoint_1.pth)
+    # Fastai busca el archivo weights_path + '.pth'
     learn.load(weights_path) 
     return learn
 
@@ -26,10 +28,11 @@ def load_clean_model(weights_path):
 learn = load_clean_model('checkpoint_1')
 
 def predict(img):
-    # PILImage.create prepara la imagen para Fastai
+    # Preparar la imagen
     img = PILImage.create(img).resize((126, 126))
     
-    # El método predict de Learner aplicará las transformaciones necesarias
+    # 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)
     
     return {categories[i]: float(probs[i]) for i in range(len(categories))}
@@ -40,7 +43,7 @@ demo = gr.Interface(
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(num_top_classes=3),
     title="Pokemon Type Classifier",
-    description="Modelo ConvNeXt Tiny cargado mediante pesos (.pth)"
+    description="Inferencia robusta con Fastai y ConvNeXt Tiny."
 )
 
 demo.launch()