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

app.py

@@ -4,44 +4,45 @@ import gradio as gr
 import timm
 import torch
 
-# 1. Lista de categorías exacta
+# 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. Reconstruir el modelo con limpieza de llaves (keys)
+# 2. Carga del modelo
 def load_pokemon_model(weights_path):
-    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 = {}
-    for k, v in state_dict.items():
-        name = k.replace('0.model.', '') if k.startswith('0.model.') else k
-        new_state_dict[name] = v
-    
-    model.load_state_dict(new_state_dict, strict=False)
-    model.eval()
-    return model
+    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)}"
 
-# 3. Preparación de Inferencia
 model = load_pokemon_model('checkpoint_1.pth')
 
+# 3. Función de predicción con captura de errores
 def predict(img):
-    # Procesamiento manual de la imagen para evitar depender de DataLoaders complejos
-    img = PILImage.create(img).resize((224, 224))
-    img_tensor = Pipeline([ToTensor(), IntToFloatTensor()])(img)
-    
-    # Añadir dimensión de batch y pasar por el modelo
-    with torch.no_grad():
-        output = model(img_tensor.unsqueeze(0))
-        probs = torch.softmax(output, dim=1)[0]
-    
-    return {categories[i]: float(probs[i]) for i in range(len(categories))}
+    try:
+        if isinstance(model, str): return {"Error de carga": model} # Si falló el paso 2
+        
+        img = PILImage.create(img).resize((224, 224))
+        
+        # Transformación manual a tensor
+        img_tensor = cast(ToTensor()(img), TensorImage)
+        img_tensor = IntToFloatTensor()(img_tensor)
+        
+        with torch.no_grad():
+            output = model(img_tensor.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)}
 
-# 4. Interfaz de Gradio
 gr.Interface(
     fn=predict,
     inputs=gr.Image(),