Update app.py

app.py

@@ -28,7 +28,7 @@ except ImportError as e:
     Sam3Model = None
     Sam3Processor = None
     root_agent = None
-    AnalysisDeps = None # Indented correctly to prevent overwriting imports [cite: 2]
+    AnalysisDeps = None
 
 # Optional: Distinctipy for better colors
 try:
@@ -256,15 +256,15 @@ def generate_overlay(image_path_str, selected_layers, layer_opacities=None, forc
 
 # --- Core Logic ---
 async def run_analysis(image_path_str, user_prompt, session_id_state):
-    waiting_df = pd.DataFrame({"Status": ["Waiting..."]})
-    empty_slider_updates = [gr.update()] * 4
+    # FIX: Use gr.skip() for updates to prevent UI jitter during streaming
+    skipped_updates = [gr.skip()] * 4
     
     if not MODEL_CACHE["loaded"]:
-        yield [], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+        yield [], None, None, [], None, gr.skip(), gr.skip(), gr.skip(), *skipped_updates
         load_models() 
 
     if not image_path_str:
-        yield [], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+        yield [], None, None, [], None, gr.skip(), gr.skip(), gr.skip(), *skipped_updates
         return
 
     # Cleanup
@@ -281,12 +281,12 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
     
     if MODEL_CACHE["model"] is None:
         error_msg = "❌ Model failed to load. Please check logs."
-        yield [{"role": "assistant", "content": error_msg}], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+        yield [{"role": "assistant", "content": error_msg}], None, None, [], None, gr.skip(), gr.skip(), gr.skip(), *skipped_updates
         return
 
     if AnalysisDeps is None:
         error_msg = "❌ Project imports failed. 'AnalysisDeps' is missing. Check your 'cellemetry' package installation."
-        yield [{"role": "assistant", "content": error_msg}], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+        yield [{"role": "assistant", "content": error_msg}], None, None, [], None, gr.skip(), gr.skip(), gr.skip(), *skipped_updates
         return
 
     try:
@@ -314,7 +314,7 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
     except Exception as e:
         error_msg = f"❌ Agent Initialization Failed: {str(e)}"
         print(error_msg)
-        yield [{"role": "assistant", "content": error_msg}], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+        yield [{"role": "assistant", "content": error_msg}], None, None, [], None, gr.skip(), gr.skip(), gr.skip(), *skipped_updates
         return
 
     image_bytes = image_path.read_bytes()
@@ -328,7 +328,6 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
 
     logs = [f"🔄 **Starting analysis** on {MODEL_CACHE['device']}..."]
     
-    # Encode spaces for Markdown
     display_path = image_path_str.replace(" ", "%20")
 
     def yield_status(log_list):
@@ -336,7 +335,8 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
         user_msg = f"![](file={display_path})\n\n{user_prompt}"
         return [{"role": "user", "content": user_msg}, {"role": "assistant", "content": full_log}]
 
-    yield yield_status(logs), session_id, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+    # FIX: Yield skips instead of updates
+    yield yield_status(logs), session_id, None, [], None, gr.skip(), gr.skip(), gr.skip(), *skipped_updates
 
     try:
         async for event in ACTIVE_RUNNER.run_async(user_id="demo_user", session_id=session.id, new_message=content):
@@ -345,7 +345,7 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
             if event.get_function_calls():
                 for fc in event.get_function_calls():
                     logs.append(f"🔧 **{author}**: Calling `{fc.name}`")
-                    yield yield_status(logs), session_id, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+                    yield yield_status(logs), session_id, None, [], None, gr.skip(), gr.skip(), gr.skip(), *skipped_updates
             
             if event.content and event.content.parts:
                 for part in event.content.parts:
@@ -360,15 +360,15 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
                                 logs[-1] = f"✅ **{author}**: {part.text}"
                             else:
                                 logs.append(f"✅ **{author}**: {part.text}")
-                        yield yield_status(logs), session_id, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+                        yield yield_status(logs), session_id, None, [], None, gr.skip(), gr.skip(), gr.skip(), *skipped_updates
 
     except Exception as e:
         logs.append(f"❌ Error: {e}")
-        yield yield_status(logs), session_id, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+        yield yield_status(logs), session_id, None, [], None, gr.skip(), gr.skip(), gr.skip(), *skipped_updates
         return
 
     logs.append("\n✅ **Analysis Complete!** Loading results...")
-    yield yield_status(logs), session_id, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+    yield yield_status(logs), session_id, None, [], None, gr.skip(), gr.skip(), gr.skip(), *skipped_updates
     
     await asyncio.sleep(0.5) 
     
@@ -385,6 +385,7 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
     final_history = [{"role": "user", "content": final_user_msg}, {"role": "assistant", "content": full_log_text}]
     slider_updates = update_opacity_sliders(layers)
     
+    # Final yield is the ONLY one with real data
     yield final_history, session_id, initial_overlay, gr.CheckboxGroup(choices=layers, value=layers), report_file, df_m, df_s, df_r, *slider_updates
 
 async def unified_chat_handler(message, history, session_id, current_img_path):
@@ -401,9 +402,10 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
         image_path = current_img_path
     
     waiting_df = pd.DataFrame({"Status": ["Waiting..."]})
-    empty_slider_updates = [gr.update()] * 4
+    # FIX: Prepare skips
+    skipped_updates = [gr.skip()] * 4
 
-    # CASE 1: INITIAL ANALYSIS (Show Loading Overlay)
+    # CASE 1: INITIAL ANALYSIS
     if image_path and (not session_id or files):
         if not user_text: user_text = "Analyze this microscopy image."
         
@@ -413,7 +415,7 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
         history.append({"role": "assistant", "content": "🔄 Starting analysis (Model loading may take a moment)..."})
 
         # Show Loading, Hide Results
-        yield history, session_id, image_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
+        yield history, session_id, image_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *skipped_updates, None, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
         
         final_result = None
         try:
@@ -423,10 +425,11 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
                 if files and len(updated_history) > 0:
                     updated_history[0] = {"role": "user", "content": f"![](file={display_path})\n\n{user_text}"}
                 
+                # Pass through the skips/data from run_analysis
                 yield (updated_history, result[1], image_path, *result[2:], None, gr.update(), gr.update(), gr.update())
         except Exception as e:
             history.append({"role": "assistant", "content": f"❌ Critical Error: {str(e)}"})
-            yield history, session_id, image_path, None, gr.CheckboxGroup(), None, waiting_df, waiting_df, waiting_df, *empty_slider_updates, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+            yield history, session_id, image_path, None, gr.CheckboxGroup(), None, gr.skip(), gr.skip(), gr.skip(), *skipped_updates, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
             return
 
         if final_result:
@@ -438,17 +441,18 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
             yield (updated_history, final_result[1], image_path, *final_result[2:], None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=True))
         return
 
-    # CASE 2: FOLLOW-UP ANALYSIS (No Loading Overlay)
+    # CASE 2: FOLLOW-UP ANALYSIS
     elif session_id and user_text:
         history.append({"role": "user", "content": user_text})
         history.append({"role": "assistant", "content": "💭 Thinking..."})
         
-        # [cite_start]Don't show loading overlay for follow-up questions [cite: 71]
-        yield history, session_id, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(), gr.update(), gr.update()
+        # Don't show loading overlay for follow-ups
+        # FIX: Send gr.skip() to all result components to prevent jitter
+        yield history, session_id, current_img_path, gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), *skipped_updates, None, gr.update(), gr.update(), gr.update()
         
         if not ACTIVE_RUNNER:
             history[-1]["content"] = "⚠️ Session expired or Agent not initialized."
-            yield history, None, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(), gr.update(), gr.update()
+            yield history, None, current_img_path, gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), *skipped_updates, None, gr.update(), gr.update(), gr.update()
             return
         
         content = types.Content(role="user", parts=[types.Part.from_text(text=user_text)])
@@ -461,10 +465,11 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
                         if hasattr(part, 'text') and part.text:
                             accumulated_response += part.text
                             history[-1]["content"] = accumulated_response
-                            yield history, session_id, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(), gr.update(), gr.update()
+                            # FIX: Keep sending gr.skip() during stream
+                            yield history, session_id, current_img_path, gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), *skipped_updates, None, gr.update(), gr.update(), gr.update()
         except Exception as e:
             history[-1]["content"] = f"❌ Error: {e}"
-            yield history, session_id, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(), gr.update(), gr.update()
+            yield history, session_id, current_img_path, gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), *skipped_updates, None, gr.update(), gr.update(), gr.update()
             return
 
         report_file, df_m, df_s, df_r = load_excel_data("") 
@@ -473,6 +478,7 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
         new_overlay = generate_overlay(current_img_path, layers, force_reload=True)
         slider_updates = update_opacity_sliders(layers)
         
+        # Final yield updates the components
         yield (
             history, 
             session_id, 
@@ -494,7 +500,7 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
             history = [{"role": "assistant", "content": "👋 Welcome! Upload a microscopy image and describe what you'd like to analyze."}]
         else:
             history.append({"role": "assistant", "content": "⚠️ Please provide a question or upload a new image."})
-        yield history, session_id, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(), gr.update(), gr.update()
+        yield history, session_id, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *skipped_updates, None, gr.update(), gr.update(), gr.update()
 
 # --- UI Layout ---
 
@@ -579,7 +585,7 @@ with gr.Blocks(title="Cellemetry Agent", css=custom_css) as demo:
                             <div style="text-align: center;">
                                 <div style="border: 8px solid #f3f3f3; border-top: 8px solid #3498db; border-radius: 50%; width: 60px; height: 60px; animation: spin 1s linear infinite; margin: 0 auto 20px;"></div>
                                 <h3 style="color: #555; margin: 0;">⚙️ Analyzing</h3>
-                                <p style="color: #888; margin-top: 10px;">Please wait while your image is processed...</p>
+                                <p style="color: #888; margin-top: 10px;">Your image is being processed...</p>
                             </div>
                             <style>@keyframes spin { 0% { transform: rotate(0deg); } 100% { transform: rotate(360deg); } }</style>
                         </div>