Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,6 +3,7 @@
 # ──── SET ENVIRONMENT VARIABLES BEFORE ANY IMPORTS ──────────────────────────────
 import os
 import tempfile
+import requests
 
 # Create dedicated cache directories
 CACHE_DIR = "/tmp/hf_cache"
@@ -14,12 +15,13 @@ os.makedirs(STREAMLIT_DIR, exist_ok=True)
 os.environ.update({
     "HOME": "/tmp",
     "XDG_CONFIG_HOME": "/tmp",
-    "STREAMLIT_HOME": STREAMLIT_DIR,  # Force Streamlit to use our temp directory
+    "STREAMLIT_HOME": STREAMLIT_DIR,
     "XDG_CACHE_HOME": CACHE_DIR,
     "HF_HOME": f"{CACHE_DIR}/huggingface",
     "TRANSFORMERS_CACHE": f"{CACHE_DIR}/transformers",
     "HF_HUB_CACHE": f"{CACHE_DIR}/huggingface_hub",
-    "HUGGINGFACE_HUB_CACHE": f"{CACHE_DIR}/huggingface_hub"
+    "HUGGINGFACE_HUB_CACHE": f"{CACHE_DIR}/huggingface_hub",
+    "HF_HUB_DISABLE_TELEMETRY": "1"  # Disable telemetry to reduce rate limiting
 })
 
 # Create all cache directories explicitly
@@ -46,11 +48,14 @@ import torchvision.transforms as T
 import streamlit as st
 from PIL import Image
 from transformers import ViTModel, T5ForConditionalGeneration, T5Tokenizer
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, HfApi
 
 # ──── MODEL DEFINITION ─────────────────────────────────────────────────────────
 MODEL_ID = "RakeshNJ12345/Chest-Radiology"
 
+# Alternative model access through proxy
+PROXY_URL = "https://hf-mirror.com"
+
 class TwoViewVisionReportModel(nn.Module):
     def __init__(self, vit: ViTModel, t5: T5ForConditionalGeneration, tokenizer: T5Tokenizer):
         super().__init__()
@@ -90,7 +95,7 @@ class TwoViewVisionReportModel(nn.Module):
         )
         return out_ids
 
-# ──── MODEL LOADING WITH ROBUST CACHE HANDLING ─────────────────────────────────
+# ──── MODEL LOADING WITH PROXY SUPPORT AND ERROR HANDLING ──────────────────────
 @st.cache_resource(show_spinner=False)
 def load_models():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -103,50 +108,106 @@ def load_models():
     ]:
         os.makedirs(path, exist_ok=True)
 
-    # Download config with explicit cache
-    cfg_path = hf_hub_download(
-        repo_id=MODEL_ID,
-        filename="config.json",
-        repo_type="model",
-        cache_dir=f"{CACHE_DIR}/huggingface_hub",
-        local_files_only=False
-    )
+    # Try to download using standard method first
+    try:
+        # Download config
+        cfg_path = hf_hub_download(
+            repo_id=MODEL_ID,
+            filename="config.json",
+            repo_type="model",
+            cache_dir=f"{CACHE_DIR}/huggingface_hub",
+            local_files_only=False
+        )
+    except Exception as e:
+        st.error(f"❌ Failed to download model via Hugging Face Hub: {str(e)}")
+        st.info("⚠️ Trying alternative download method...")
+        
+        # Use proxy mirror
+        cfg_path = f"{CACHE_DIR}/huggingface_hub/config.json"
+        api = HfApi(endpoint=PROXY_URL)
+        api.hf_hub_download(
+            repo_id=MODEL_ID,
+            filename="config.json",
+            repo_type="model",
+            cache_dir=f"{CACHE_DIR}/huggingface_hub",
+            local_files_only=False
+        )
+
     cfg = json.load(open(cfg_path, "r"))
 
     # Load models with explicit cache directories
-    vit = ViTModel.from_pretrained(
-        "google/vit-base-patch16-224",
-        ignore_mismatched_sizes=True,
-        cache_dir=f"{CACHE_DIR}/transformers"
-    ).to(device)
+    try:
+        vit = ViTModel.from_pretrained(
+            "google/vit-base-patch16-224",
+            ignore_mismatched_sizes=True,
+            cache_dir=f"{CACHE_DIR}/transformers"
+        ).to(device)
+    except:
+        # Use proxy if standard download fails
+        vit = ViTModel.from_pretrained(
+            "google/vit-base-patch16-224",
+            ignore_mismatched_sizes=True,
+            cache_dir=f"{CACHE_DIR}/transformers",
+            mirror=PROXY_URL
+        ).to(device)
     
-    t5 = T5ForConditionalGeneration.from_pretrained(
-        "t5-base",
-        cache_dir=f"{CACHE_DIR}/transformers"
-    ).to(device)
+    try:
+        t5 = T5ForConditionalGeneration.from_pretrained(
+            "t5-base",
+            cache_dir=f"{CACHE_DIR}/transformers"
+        ).to(device)
+    except:
+        t5 = T5ForConditionalGeneration.from_pretrained(
+            "t5-base",
+            cache_dir=f"{CACHE_DIR}/transformers",
+            mirror=PROXY_URL
+        ).to(device)
     
-    tok = T5Tokenizer.from_pretrained(
-        MODEL_ID,
-        cache_dir=f"{CACHE_DIR}/transformers"
-    )
+    try:
+        tok = T5Tokenizer.from_pretrained(
+            MODEL_ID,
+            cache_dir=f"{CACHE_DIR}/transformers"
+        )
+    except:
+        tok = T5Tokenizer.from_pretrained(
+            MODEL_ID,
+            cache_dir=f"{CACHE_DIR}/transformers",
+            mirror=PROXY_URL
+        )
 
     # Load combined model
     model = TwoViewVisionReportModel(vit, t5, tok).to(device)
     
-    ckpt_path = hf_hub_download(
-        repo_id=MODEL_ID,
-        filename="pytorch_model.bin",
-        repo_type="model",
-        cache_dir=f"{CACHE_DIR}/huggingface_hub",
-        local_files_only=False
-    )
+    try:
+        ckpt_path = hf_hub_download(
+            repo_id=MODEL_ID,
+            filename="pytorch_model.bin",
+            repo_type="model",
+            cache_dir=f"{CACHE_DIR}/huggingface_hub",
+            local_files_only=False
+        )
+    except:
+        # Use proxy mirror for model weights
+        api = HfApi(endpoint=PROXY_URL)
+        ckpt_path = api.hf_hub_download(
+            repo_id=MODEL_ID,
+            filename="pytorch_model.bin",
+            repo_type="model",
+            cache_dir=f"{CACHE_DIR}/huggingface_hub",
+            local_files_only=False
+        )
     
     state = torch.load(ckpt_path, map_location=device)
     model.load_state_dict(state)
     return device, model, tok
 
-# ──── APP INTERFACE ───────────────────────────────────────────────────────────
-device, model, tokenizer = load_models()
+# ──── APP INTERFACE WITH ERROR HANDLING ───────────────────────────────────────
+try:
+    device, model, tokenizer = load_models()
+except Exception as e:
+    st.error(f"🚨 Critical Error: Failed to load models. {str(e)}")
+    st.stop()
+
 transform = T.Compose([
     T.Resize((224, 224)),
     T.ToTensor(),
@@ -161,8 +222,15 @@ st.markdown("<p style='text-align:center;'>Upload a chest X-ray and click Genera
 if "img" not in st.session_state:
     uploaded = st.file_uploader("📤 Upload X-ray (PNG/JPG)", type=["png", "jpg", "jpeg"])
     if uploaded:
-        st.session_state.img = uploaded
-        st.experimental_rerun()
+        try:
+            # Validate image
+            img = Image.open(uploaded).convert("RGB")
+            img.verify()  # Check if image is valid
+            st.session_state.img = uploaded
+            st.experimental_rerun()
+        except Exception as e:
+            st.error(f"❌ Invalid image file: {str(e)}")
+            st.stop()
     else:
         st.stop()
 
@@ -174,14 +242,27 @@ col1, col2 = st.columns(2)
 with col1:
     if st.button("▶️ Generate Report", use_container_width=True):
         with st.spinner("Analyzing X-ray..."):
-            px = transform(img).unsqueeze(0).to(device)
-            out_ids = model.generate(px, max_length=128)
-            report = tokenizer.decode(out_ids[0], skip_special_tokens=True)
-        
-        st.subheader("📝 AI-Generated Report")
-        st.success(report)
+            try:
+                px = transform(img).unsqueeze(0).to(device)
+                out_ids = model.generate(px, max_length=128)
+                report = tokenizer.decode(out_ids[0], skip_special_tokens=True)
+                
+                st.subheader("📝 AI-Generated Report")
+                st.success(report)
+            except Exception as e:
+                st.error(f"❌ Analysis failed: {str(e)}")
+                st.info("Please try with a different image or try again later")
 
 with col2:
     if st.button("⬅️ Upload Another", use_container_width=True):
         del st.session_state.img
-        st.experimental_rerun()
+        st.experimental_rerun()
+
+# Add footer with troubleshooting
+st.markdown("---")
+st.markdown("""
+**Troubleshooting Tips:**
+- If model download fails, wait 5 minutes and refresh
+- Use standard chest X-ray images in PNG or JPG format
+- For persistent errors, contact support@example.com
+""")