Update streamlit_app.py

streamlit_app.py

@@ -25,12 +25,25 @@ except:
 
 @st.cache_resource
 def load_model():
+    import os
+    import urllib.request
+
+    MODEL_URL = "https://huggingface.co/your-username/your-repo/resolve/main/butterfly_classifier.pth"
+    MODEL_PATH = "butterfly_classifier.pth"
+
+    # Download if not already present
+    if not os.path.exists(MODEL_PATH):
+        st.info("Downloading model...")
+        urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
+
+    # Load the model
     model = models.resnet18(pretrained=False)
     model.fc = torch.nn.Linear(model.fc.in_features, len(class_names))
-    model.load_state_dict(torch.load("butterfly_classifier.pth", map_location="cpu"))
+    model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
     model.eval()
     return model
 
+
 model = load_model()
 
 transform = transforms.Compose([