Forzando sobrescritura de app.py y carga de pesos .pth

app.py

@@ -21,24 +21,24 @@ model = load_pokemon_model('checkpoint_1.pth')
 
 def predict(img):
     try:
-        # 1. Preparar la imagen al tamaño exacto de tu entrenamiento (126)
+        # 1. Redimensionar
         img = PILImage.create(img).resize((126, 126))
         
-        # 2. Transformar a Tensor y Normalizar (Estándar de ConvNeXt/ImageNet)
-        # Esto soluciona el error de los canales (channels)
-        timg = TensorImage(image2tensor(img))
-        timg = IntToFloatTensor()(timg)
-        timg = Normalize.from_stats(*imagenet_stats)(timg)
+        # 2. Convertir a tensor y normalizar (forma directa de PyTorch)
+        # Usamos transforms estándar para asegurar las 4 dimensiones [1, 3, 126, 126]
+        timg = Pipeline([ToTensor(), IntToFloatTensor(), Normalize.from_stats(*imagenet_stats)])(img)
+        
+        # 3. Predicción (timg ya sale con el batch dimension si se usa Pipeline de fastai adecuadamente)
+        # Pero para estar seguros de que tiene 4D y no 5D:
+        if timg.ndim == 3: timg = timg.unsqueeze(0)
+        if timg.ndim == 5: timg = timg.squeeze(0) # Esto elimina la dimensión extra si existiera
         
-        # 3. Predicción
         with torch.no_grad():
-            output = model(timg.unsqueeze(0))
+            output = model(timg)
             probs = torch.softmax(output, dim=1)[0]
         
         return {categories[i]: float(probs[i]) for i in range(len(categories))}
     except Exception as e:
-        # Devolvemos un string que Gradio pueda manejar en caso de error crítico
-        print(f"Error en predicción: {e}")
         return {f"Error: {str(e)}": 0.0}
 
 gr.Interface(