Update src/streamlit_app.py

src/streamlit_app.py

@@ -24,19 +24,31 @@ from transformers import ViTModel, T5ForConditionalGeneration, T5Tokenizer
 # ─── point at your 1.2 GB model repo, NOT this Space —──────────────────────────
 HF_MODEL_ID = "RakeshNJ12345/Chest-Radiology"
 
+@st.cache_resource(show_spinner=False)
 @st.cache_resource(show_spinner=False)
 def load_models():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # 1) Vision trunk (fine‐tuned ViT weights under `vit/` in your model repo)
-    vit = ViTModel.from_pretrained(f"{HF_MODEL_ID}/vit").to(device)
-
-    # 2) T5 + tokenizer (your fine‐tuned report generator)
-    t5  = T5ForConditionalGeneration.from_pretrained(HF_MODEL_ID).to(device)
-    tok = T5Tokenizer.from_pretrained(HF_MODEL_ID)
+    # 1) Frozen ViT pulled from the "vit/" folder of your model repo
+    vit = ViTModel.from_pretrained(
+        HF_MODEL_ID,
+        subfolder="vit",
+        local_files_only=True,
+    ).to(device)
+
+    # 2) Fine-tuned T5 & tokenizer at the root of that same repo
+    t5  = T5ForConditionalGeneration.from_pretrained(
+        HF_MODEL_ID,
+        local_files_only=True,
+    ).to(device)
+    tok = T5Tokenizer.from_pretrained(
+        HF_MODEL_ID,
+        local_files_only=True,
+    )
 
     return device, vit, t5, tok
 
+
 device, vit, t5, tokenizer = load_models()
 
 # ─── preprocessing for ViT —────────────────────────────────────────────────────