Update app.py

app.py

@@ -4,6 +4,7 @@ from torch.utils.data import Dataset, DataLoader
 import torchvision
 from PIL import Image
 import gradio as gr
+from huggingface_hub import hf_hub_download
 
 # Configure device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -14,7 +15,11 @@ transform = torchvision.transforms.Compose([
     torchvision.transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.225, 0.225, 0.225])
 ])
 
-model = torch.jit.load('model_smoke.pt')
+REPO_ID = "Raaniel/model-smoke"
+MODEL_FILE_NAME = "model_smoke.pt"
+checkpoint_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILE_NAME)
+
+model = torch.jit.load(checkpoint_path)
 model = model.to(device)
 
 classes = ["chmury", 'inne', "dym"]