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

app.py

@@ -4,44 +4,42 @@ import gradio as gr
 import timm
 import torch
 
-# 1. Lista de categorías
 categories = ['Bug', 'Dark', 'Dragon', 'Electric', 'Fairy', 'Fighting', 
               'Fire', 'Flying', 'Ghost', 'Grass', 'Ground', 'Ice', 
               'Normal', 'Poison', 'Psychic', 'Rock', 'Steel', 'Water']
 
-# 2. Carga del modelo
 def load_pokemon_model(weights_path):
-    try:
-        model = timm.create_model('convnext_tiny', pretrained=False, num_classes=len(categories))
-        state_dict = torch.load(weights_path, map_location='cpu', weights_only=False)
-        if 'model' in state_dict: state_dict = state_dict['model']
-        new_state_dict = {k.replace('0.model.', ''): v for k, v in state_dict.items()}
-        model.load_state_dict(new_state_dict, strict=False)
-        model.eval()
-        return model
-    except Exception as e:
-        return f"Error cargando pesos: {str(e)}"
+    model = timm.create_model('convnext_tiny', pretrained=False, num_classes=len(categories))
+    state_dict = torch.load(weights_path, map_location='cpu', weights_only=False)
+    if 'model' in state_dict: state_dict = state_dict['model']
+    new_state_dict = {k.replace('0.model.', ''): v for k, v in state_dict.items()}
+    model.load_state_dict(new_state_dict, strict=False)
+    model.eval()
+    return model
 
 model = load_pokemon_model('checkpoint_1.pth')
 
-# 3. Función de predicción con captura de errores
 def predict(img):
     try:
-        if isinstance(model, str): return {"Error de carga": model} # Si falló el paso 2
-        
-        img = PILImage.create(img).resize((224, 224))
+        # 1. Preparar la imagen al tamaño exacto de tu entrenamiento (126)
+        img = PILImage.create(img).resize((126, 126))
         
-        # Transformación manual a tensor
-        img_tensor = cast(ToTensor()(img), TensorImage)
-        img_tensor = IntToFloatTensor()(img_tensor)
+        # 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)
         
+        # 3. Predicción
         with torch.no_grad():
-            output = model(img_tensor.unsqueeze(0))
+            output = model(timg.unsqueeze(0))
             probs = torch.softmax(output, dim=1)[0]
         
         return {categories[i]: float(probs[i]) for i in range(len(categories))}
     except Exception as e:
-        return {"Error en predicción": str(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(
     fn=predict,