Usando from_pretrained_fastai para carga limpia

app.py

@@ -1,52 +1,39 @@
 
-from fastai.vision.all import *
 import gradio as gr
-import timm
-import torch
+from fastai.vision.all import *
+from huggingface_hub import from_pretrained_fastai
+import torch, os
+
+# Optimizaciones de CPU para el servidor
+os.environ.setdefault("OMP_NUM_THREADS", "1")
+torch.set_num_threads(1)
 
-categories = ['Bug', 'Dark', 'Dragon', 'Electric', 'Fairy', 'Fighting', 
-              'Fire', 'Flying', 'Ghost', 'Grass', 'Ground', 'Ice', 
-              'Normal', 'Poison', 'Psychic', 'Rock', 'Steel', 'Water']
+# --- REEMPLAZA CON TU REPO_ID ---
+# Ejemplo: "Edupy/pokemon-1class-classifier-26"
+learn = from_pretrained_fastai("Edupy/pokemon-1class-classifier-26")
 
-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 = {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
+try:
+    learn.to_fp32()
+except:
+    pass
 
-model = load_pokemon_model('checkpoint_1.pth')
+labels = learn.dls.vocab
 
 def predict(img):
-    try:
-        # 1. Procesar imagen PIL
-        img = PILImage.create(img).resize((126, 126))
-        
-        # 2. Convertir a tensor manualmente para tener control total
-        # image2tensor devuelve un tensor de [3, 126, 126]
-        timg = image2tensor(img)
-        timg = timg.float()/255.0 # Normalizar a 0-1
-        
-        # 3. Aplicar normalización de ImageNet (Media y Desviación estándar)
-        stats = imagenet_stats
-        timg = (timg - tensor(stats[0])[:,None,None]) / tensor(stats[1])[:,None,None]
-        
-        # 4. Añadir dimensión de batch -> Resultado: [1, 3, 126, 126]
-        batch_img = timg.unsqueeze(0)
-        
-        with torch.no_grad():
-            output = model(batch_img)
-            probs = torch.softmax(output, dim=1)[0]
-        
-        return {categories[i]: float(probs[i]) for i in range(len(categories))}
-    except Exception as e:
-        return {f"Error: {str(e)}": 0.0}
-
-gr.Interface(
+    img = PILImage.create(img)
+    pred, pred_idx, probs = learn.predict(img)
+    return {labels[i]: float(probs[i]) for i in range(len(labels))}
+
+title = "Detector de Tipos Pokémon"
+description = "Modelo profesional cargado desde el Hub para identificar tipos elementales."
+
+demo = gr.Interface(
     fn=predict,
-    inputs=gr.Image(),
+    inputs=gr.Image(type="pil"),
     outputs=gr.Label(num_top_classes=3),
-    title="Detector de Tipos Pokémon"
-).launch()
+    title=title,
+    description=description
+)
+
+# Configuración de cola para evitar que el Space se cuelgue con muchas peticiones
+demo.queue(max_size=8).launch(show_error=True, debug=True)