Update app.py

app.py

@@ -6,7 +6,13 @@ import torch
 
 
 CHECKPOINT = "g8a9/vit-geppetto-captioning"
-model = VisionEncoderDecoderModel.from_pretrained(CHECKPOINT)
+
+@st.cache
+def get_model():
+    model = VisionEncoderDecoderModel.from_pretrained(CHECKPOINT)
+    return model
+
+
 feature_extractor = AutoFeatureExtractor.from_pretrained(CHECKPOINT)
 tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
 
@@ -42,6 +48,7 @@ elif gen_mode == "sampling":
 def generate_caption(url):
     image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
     inputs = feature_extractor(image, return_tensors="pt")
+    model = get_model()
     generated_ids = model.generate(
         inputs["pixel_values"],
         max_length=20,