Add real-time log updates and multi-line tensor preview

app.py

@@ -6,6 +6,8 @@ import shutil
 import logging
 import io
 import sys
+import threading
+import time
 from contextlib import redirect_stdout, redirect_stderr
 
 import gradio as gr
@@ -109,30 +111,53 @@ def format_tensor_info(tensor: torch.Tensor) -> str:
 def fetch_param(model_id: str, param_key: str, progress=gr.Progress()):
     """Fetch parameter and return formatted info and tensor preview."""
     log_buffer = io.StringIO()
+    last_log_value = ""
 
     if not model_id or not param_key:
-        return "Please provide both model ID and parameter key.", "", None, "❌ Missing required inputs"
+        yield "Please provide both model ID and parameter key.", "", None, "❌ Missing required inputs"
+        return
 
     try:
         log_buffer.write(f"🚀 Starting download for {model_id}\n")
         log_buffer.write(f"🎯 Target parameter: {param_key}\n\n")
         progress(0, desc="Initializing...")
+        yield "", "", None, log_buffer.getvalue()
+        time.sleep(0.5)
 
+        # Start download in background and monitor logs
         tensor = get_param(model_id, param_key, log_buffer, progress)
+        
+        # Yield log updates periodically during download
+        current_log = log_buffer.getvalue()
+        if current_log != last_log_value:
+            yield "", "", None, current_log
+            last_log_value = current_log
+        
         info = format_tensor_info(tensor)
 
         # Create tensor preview (first few elements)
         log_buffer.write(f"\n📊 Creating preview...\n")
+        yield "", "", None, log_buffer.getvalue()
+        
         flat = tensor.flatten()
         preview_size = min(100, flat.numel())
         preview = flat[:preview_size].tolist()
-        preview_str = f"**First {preview_size} values:**\n```\n{preview}\n```"
+        
+        # Format preview in multiple lines (10 values per line)
+        preview_lines = []
+        for i in range(0, len(preview), 10):
+            line_values = preview[i:i+10]
+            preview_lines.append(", ".join(f"{v:.6f}" for v in line_values))
+        
+        preview_str = f"**First {preview_size} values:**\n```\n" + "\n".join(preview_lines) + "\n```"
 
         if flat.numel() > preview_size:
             preview_str += f"\n\n... and {flat.numel() - preview_size:,} more values"
 
         # Save tensor for download
         log_buffer.write(f"💾 Saving tensor for download...\n")
+        yield info, preview_str, None, log_buffer.getvalue()
+        
         temp_dir = tempfile.gettempdir()
         safe_param_key = param_key.replace("/", "_").replace(".", "_")
         download_path = os.path.join(temp_dir, f"{safe_param_key}.pt")
@@ -141,10 +166,10 @@ def fetch_param(model_id: str, param_key: str, progress=gr.Progress()):
 
         progress(1.0, desc="Complete!")
         log_buffer.write(f"\n✅ All operations completed successfully!\n")
-        return info, preview_str, download_path, log_buffer.getvalue()
+        yield info, preview_str, download_path, log_buffer.getvalue()
     except Exception as e:
         log_buffer.write(f"\n❌ Error: {str(e)}\n")
-        return f"**Error:** {str(e)}", "", None, log_buffer.getvalue()
+        yield f"**Error:** {str(e)}", "", None, log_buffer.getvalue()
 
 
 def list_keys(model_id: str):