clean model loaded from hf

.gitattributes

@@ -1 +0,0 @@
-*.pt filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -13,8 +13,8 @@ venv/
 .DS_Store
 
 # PyTorch checkpoints
-# *.pt
-#!checkpoints/transformer_mnist.pt  # allow the model to be pushed to HF Space
+checkpoints/
+*.pt
 
 # Gradio session files
 gradio_cached_examples/

gradio_ui/gradio_app.py

@@ -5,6 +5,7 @@ import gradio as gr
 import torch
 import numpy as np
 from PIL import Image, ImageDraw
+from huggingface_hub import hf_hub_download
 
 from model.model import ImageToDigitTransformer
 from utils.tokenizer import START, FINISH, decode
@@ -12,8 +13,10 @@ from gradio_ui.preprocess import preprocess_canvases
 
 # Load model
 device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+
+model_path = hf_hub_download(repo_id="nico-x/vit-transformer-mnist", filename="transformer_mnist.pt")
 model = ImageToDigitTransformer(vocab_size=13).to(device)
-model.load_state_dict(torch.load("checkpoints/transformer_mnist.pt", map_location=device))
+model.load_state_dict(torch.load(model_path, map_location=device))
 model.eval()
 
 def split_into_quadrants(image):