feat: Add download status polling to UI and Firefox extension

firefox-extension/content.js

@@ -48,6 +48,26 @@ browser.runtime.onMessage.addListener(async (request, sender, sendResponse) => {
     const checkedLayers = [4, 5, 6, 7, 8, 9, 10, 11, 12, 13];
 
     try {
+      // Start polling status
+      let isFetching = true;
+      const pollStatus = async () => {
+        while (isFetching) {
+          try {
+            const statusRes = await fetch(`${apiUrl}/status`);
+            if (statusRes.ok) {
+              const statusData = await statusRes.json();
+              if (statusData[modelVersion] === "downloading") {
+                showToast(`Downloading Gemma 4 (${modelVersion})... this may take a few minutes.`, 0);
+              }
+            }
+          } catch (e) {
+            // ignore network errors for status polling
+          }
+          await new Promise(r => setTimeout(r, 2000));
+        }
+      };
+      pollStatus();
+
       // Connect to the configured API Space
       const response = await fetch(`${apiUrl}/analyze/${modelVersion}`, {
         method: 'POST',
@@ -59,6 +79,7 @@ browser.runtime.onMessage.addListener(async (request, sender, sendResponse) => {
           layers: checkedLayers 
         })
       });
+      isFetching = false;
 
       if (!response.ok) throw new Error('API error');
       
@@ -159,9 +180,32 @@ async function processEntirePage() {
   const batchSize = 3;
   let processedCount = 0;
 
+  // Polling logic for first request
+  let isFetchingStatus = true;
+  const pollStatus = async () => {
+    while (isFetchingStatus) {
+      try {
+        const statusRes = await fetch(`${apiUrl}/status`);
+        if (statusRes.ok) {
+          const statusData = await statusRes.json();
+          if (statusData[modelVersion] === "downloading") {
+            showToast(`Downloading Gemma 4 (${modelVersion})... this may take a few minutes.`, 0);
+          }
+        }
+      } catch (e) {}
+      await new Promise(r => setTimeout(r, 2000));
+    }
+  };
+  pollStatus();
+
   for (let i = 0; i < nodesToProcess.length; i += batchSize) {
     const batch = nodesToProcess.slice(i, i + batchSize);
-    showToast(`FlowRead analyzing page (${processedCount}/${nodesToProcess.length} blocks)...`, 0);
+    
+    // Only show analyzing text if not downloading
+    const statusResTemp = await fetch(`${apiUrl}/status`).catch(() => null);
+    if (!statusResTemp || !statusResTemp.ok || (await statusResTemp.json())[modelVersion] !== "downloading") {
+       showToast(`FlowRead analyzing page (${processedCount}/${nodesToProcess.length} blocks)...`, 0);
+    }
 
     await Promise.all(batch.map(async (node) => {
       const text = node.nodeValue;
@@ -206,6 +250,7 @@ async function processEntirePage() {
     processedCount += batch.length;
   }
 
+  isFetchingStatus = false;
   showToast(`Done! Analyzed ${processedCount} blocks.`, 2000);
 }
 
@@ -223,6 +268,7 @@ async function updateExisting(newSettings) {
 
   let reFetchCount = 0;
   let rerenderCount = 0;
+  let isFetchingStatus = false;
 
   for (const container of containers) {
     const oldPreprompt = container.dataset.preprompt || "";
@@ -234,6 +280,21 @@ async function updateExisting(newSettings) {
     if (oldPreprompt !== preprompt || oldMode !== saliencyMode || oldModelVersion !== modelVersion) {
       if (reFetchCount === 0) {
          showToast("Updating FlowRead elements with new settings...", 0);
+         isFetchingStatus = true;
+         (async () => {
+           while (isFetchingStatus) {
+             try {
+               const statusRes = await fetch(`${apiUrl}/status`);
+               if (statusRes.ok) {
+                 const statusData = await statusRes.json();
+                 if (statusData[modelVersion] === "downloading") {
+                   showToast(`Downloading Gemma 4 (${modelVersion})... this may take a few minutes.`, 0);
+                 }
+               }
+             } catch (e) {}
+             await new Promise(r => setTimeout(r, 2000));
+           }
+         })();
       }
       try {
         const response = await fetch(`${apiUrl}/analyze/${modelVersion}`, {
@@ -274,6 +335,8 @@ async function updateExisting(newSettings) {
     }
   }
 
+  isFetchingStatus = false;
+
   if (reFetchCount > 0) {
     showToast(`Updated ${reFetchCount} blocks with new AI intent!`, 2000);
   } else if (rerenderCount > 0) {

main.py

@@ -236,6 +236,10 @@ def get_study_stats():
 # --- Saliency API (Existing) ---
 models = {}
 tokenizers = {}
+model_status = {
+    "2b": "unloaded",
+    "27b-4a": "unloaded"
+}
 
 device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
 hf_token = os.environ.get("HF_TOKEN")
@@ -245,6 +249,7 @@ def load_model(model_name: str):
         return models[model_name], tokenizers[model_name]
         
     print(f"Loading {model_name} on {device}...")
+    model_status[model_name] = "downloading"
     try:
         if model_name == "27b-4a":
             # Use Gemma 4 27B in 4-bit (requires CUDA)
@@ -279,9 +284,11 @@ def load_model(model_name: str):
         print(f"Model {model_name} loaded successfully.")
         models[model_name] = model
         tokenizers[model_name] = tokenizer
+        model_status[model_name] = "loaded"
         return model, tokenizer
     except Exception as e:
         print(f"Error loading model {model_name}: {e}")
+        model_status[model_name] = "error"
         raise e
 
 # Pre-load default 2b
@@ -290,6 +297,10 @@ try:
 except:
     print("Could not preload 2b model.")
 
+@app.get("/status")
+def get_model_status():
+    return model_status
+
 class TextRequest(BaseModel):
     text: str
     layers: Optional[List[int]] = None  # List of layer indices to average
@@ -297,11 +308,11 @@ class TextRequest(BaseModel):
     saliency_mode: str = "local" # "local" or "global"
 
 @app.post("/analyze")
-async def analyze_text_legacy(request: TextRequest):
-    return await analyze_text_model("2b", request)
+def analyze_text_legacy(request: TextRequest):
+    return analyze_text_model("2b", request)
 
 @app.post("/analyze/{model_name}")
-async def analyze_text_model(model_name: str, request: TextRequest):
+def analyze_text_model(model_name: str, request: TextRequest):
     text = request.text
     preprompt = request.preprompt.strip()
     

static/index.html

@@ -463,6 +463,7 @@
             const text = inputArea.value.trim();
             const preprompt = prepromptInput.value.trim();
             const saliencyMode = saliencyModeInput.value;
+            const modelVersion = "2b";
             
             const checkedLayers = Array.from(document.querySelectorAll('.layer-cb:checked')).map(cb => parseInt(cb.value));
             
@@ -474,14 +475,34 @@
 
             analyzeBtn.disabled = true;
             loading.style.display = 'block';
+            loading.textContent = 'Analyzing text with Gemma 4...';
             resultContainer.innerHTML = '';
 
+            let isFetching = true;
+            const pollStatus = async () => {
+                while(isFetching) {
+                    try {
+                        const statusRes = await fetch('/status');
+                        if (statusRes.ok) {
+                            const statusData = await statusRes.json();
+                            if (statusData[modelVersion] === "downloading") {
+                                loading.textContent = `Downloading Gemma 4 (${modelVersion}) Model... this takes a few minutes.`;
+                            }
+                        }
+                    } catch(e) {}
+                    await new Promise(r => setTimeout(r, 2000));
+                }
+            };
+            pollStatus();
+
             try {
-                const response = await fetch('/analyze', {
+                const response = await fetch(`/analyze/${modelVersion}`, {
                     method: 'POST',
                     headers: { 'Content-Type': 'application/json' },
                     body: JSON.stringify({ text, preprompt, layers: checkedLayers, saliency_mode: saliencyMode })
                 });
+                
+                isFetching = false;
 
                 if (!response.ok) throw new Error('Network response was not ok');
 
@@ -783,7 +804,11 @@
 
             } catch (err) {
                 console.error(err);
-                alert("Error loading stats.");
+                alert("Failed to analyze text.");
+                isFetching = false;
+            } finally {
+                analyzeBtn.disabled = false;
+                loading.style.display = 'none';
             }
         }
     </script>