Update app.py

app.py

@@ -4,16 +4,20 @@ from torch import nn
 import numpy as np
 import gradio as gr
 import PIL
+from huggingface_hub import from_pretrained_fastai
 
 # Instantiate classification model
 from fastai.vision.all import *
-model_multi = load_learner('vit_tiny_patch16.pkl')
+#model_multi = load_learner('vit_tiny_patch16.pkl')
+repo_id = "zklee98/vit_tiny_patch16_multi"
+model_multi = from_pretrained_fastai(repo_id)
 
 def binary_label(path):
     return 'No-anomaly' if (parent_label(path) == 'No-Anomaly') else 'Anomaly'
 
-model_binary = load_learner('vit_tiny_patch16_binary.pkl')
-
+#model_binary = load_learner('vit_tiny_patch16_binary.pkl')
+model_binary = from_pretrained_fastai(repo_id='zklee98/vit_tiny_patch16_binary')
+    
 # Instantiate segmentation model
 from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
 from torchvision.transforms import Grayscale