Update app.py

app.py

@@ -1,53 +1,20 @@
 import gradio as gr
-import traceback
 from fastai.vision.all import *
 from huggingface_hub import from_pretrained_fastai
 import torch, os
-import torchvision.transforms as T
 
 os.environ.setdefault("OMP_NUM_THREADS", "1")
 torch.set_num_threads(1)
 
 learn = from_pretrained_fastai("Pablogps/castle-classifier-25")
-try:
-    learn.to_fp32()
-except Exception:
-    pass
-
-for dl in learn.dls.loaders:
-    try: dl.num_workers = 0
-    except: pass
-
+try: learn.to_fp32()
+except: pass
 labels = learn.dls.vocab
 
-
 def predict(img):
-    try:
-        print("Starting prediction...", flush=True)
-        img = PILImage.create(img)
-
-        # 2) Torchvision preprocessing:
-        #    Resize shorter side to 256, center-crop 224, ToTensor, ImageNet normalize
-        tfm = T.Compose([
-            T.Resize(256),
-            T.CenterCrop(224),
-            T.ToTensor(),  # -> [C,H,W] in [0,1]
-            T.Normalize(mean=[0.485, 0.456, 0.406],
-                        std=[0.229, 0.224, 0.225]),
-        ])
-        x = tfm(img).unsqueeze(0)  # [1,3,224,224]
-
-        # 3) Forward pass
-        learn.model.eval()
-        logits = learn.model(x)                 # [1, C]
-        probs = torch.softmax(logits, dim=1)[0] # [C]
-        
-        # pred, pred_idx, probs = learn.predict(img)
-        return {labels[i]: float(probs[i]) for i in range(len(labels))}
-    except Exception as e:
-        tb = traceback.format_exc()
-        print(tb, flush=True)          # goes to runtime logs even if you can't see them
-        raise gr.Error(tb)             # shows full traceback to you in the UI
+    img = PILImage.create(img)          # same flow as before
+    pred, pred_idx, probs = learn.predict(img)
+    return {labels[i]: float(probs[i]) for i in range(len(labels))}
 
 title = "Bad castle predictor"
 description = "A bad model that tries to identify the type of castle."