Update app.py

app.py

@@ -4,6 +4,8 @@ import yaml
 import shutil
 import subprocess
 import sys
+import queue
+import threading
 from datetime import datetime
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Iterable
@@ -277,51 +279,87 @@ def run_mineru_parsing_and_dag_gen():
         error_log = full_log + f"\n[Global Exception] Exception occurred:\n{str(e)}"
         yield "❌ Execution Exception", get_debug_info(), error_log
 
+
 def run_final_generation(task_type="all"):
+    """
+    使用队列和多线程实时获取子进程的流式输出，并通过 yield 推送到前端界面。
+    """
     if not os.path.exists(OUTPUT_DIR):
-        return "❌ Please run the parsing step first", get_debug_info(), "No output folder found.", None
+        yield "❌ Please run the parsing step first", get_debug_info(), "No output folder found.", gr.update(visible=False)
+        return
 
     scripts_to_run = []
     if task_type == "ppt": scripts_to_run = ["gen_ppt.py"]
     elif task_type == "poster": scripts_to_run = ["gen_poster.py"]
     elif task_type == "pr": scripts_to_run = ["gen_pr.py"]
     elif task_type == "all": scripts_to_run = ["gen_ppt.py", "gen_poster.py", "gen_pr.py"]
-    else: return "❌ Unknown task type", get_debug_info(), "Invalid task_type.", None
+    else: 
+        yield "❌ Unknown task type", get_debug_info(), "Invalid task_type.", gr.update(visible=False)
+        return
 
     full_log = f"🚀 Preparing to start {len(scripts_to_run)} tasks...\n"
-    success = True
+    yield f"⏳ Starting {task_type.upper()} generation...", get_debug_info(), full_log, gr.update(visible=False)
 
-    def execute_script(script):
-        command = [sys.executable, script]
-        return script, subprocess.run(command, capture_output=True, text=True, timeout=600)
+    q = queue.Queue()
+    processes = []
+    
+    # 将标准输出推送到队列的工作线程
+    def enqueue_output(out, script_name):
+        for line in iter(out.readline, ''):
+            q.put((script_name, line))
+        out.close()
 
     try:
-        with ThreadPoolExecutor(max_workers=len(scripts_to_run)) as executor:
-            future_to_script = {executor.submit(execute_script, s): s for s in scripts_to_run}
-            for future in as_completed(future_to_script):
-                script_name = future_to_script[future]
-                try:
-                    _, result = future.result()
-                    full_log += f"\n================ ✅ Execution Complete: {script_name} ================\n"
-                    full_log += f"--- STDOUT ---\n{result.stdout}\n\n--- STDERR ---\n{result.stderr}\n"
-                    if result.returncode != 0:
-                        success = False
-                        full_log += f"❌ [Error] {script_name} returned non-zero exit code (Exit Code: {result.returncode})\n"
-                except subprocess.TimeoutExpired as e:
-                    success = False
-                    full_log += f"\n================ ❌ Task Timeout: {script_name} ================\n{str(e)}\n"
-                except Exception as e:
-                    success = False
-                    full_log += f"\n================ ❌ Task Exception: {script_name} ================\n{str(e)}\n"
-
-        if not success: return f"❌ {task_type.upper()} contains failed tasks, please check logs", get_debug_info(), full_log, None
+        # 并发启动所有脚本
+        for script in scripts_to_run:
+            p = subprocess.Popen(
+                [sys.executable, script], 
+                stdout=subprocess.PIPE, 
+                stderr=subprocess.STDOUT, 
+                text=True, 
+                bufsize=1
+            )
+            processes.append((script, p))
+            t = threading.Thread(target=enqueue_output, args=(p.stdout, script))
+            t.daemon = True
+            t.start()
+
+        # 从队列中实时读取日志并更新 UI
+        active_processes = len(processes)
+        while active_processes > 0 or not q.empty():
+            try:
+                # 设定 timeout，这样即便没有日志产生也能循环检查进程是否结束
+                script_name, line = q.get(timeout=0.1)
+                full_log += f"[{script_name}] {line}"
+                yield f"⏳ Generating {task_type.upper()}...", get_debug_info(), full_log, gr.update(visible=False)
+            except queue.Empty:
+                active_processes = sum(1 for _, p in processes if p.poll() is None)
+
+        # 检查最终执行结果
+        success = True
+        for script, p in processes:
+            if p.returncode != 0:
+                success = False
+                full_log += f"\n❌ [Error] {script} returned non-zero exit code (Exit Code: {p.returncode})\n"
+
+        if not success: 
+            yield f"❌ {task_type.upper()} contains failed tasks, please check logs", get_debug_info(), full_log, gr.update(visible=False)
+            return
 
+        # 全部成功后进行压缩
+        full_log += "\n📦 Zipping output directory...\n"
+        yield f"��� Zipping outputs...", get_debug_info(), full_log, gr.update(visible=False)
+        
         zip_base_name = ZIP_OUTPUT_PATH.replace(".zip", "")
         shutil.make_archive(zip_base_name, 'zip', OUTPUT_DIR)
-        return f"✅ {task_type.upper()} generated and zipped successfully", get_debug_info(), full_log, ZIP_OUTPUT_PATH
+        
+        full_log += "✅ All tasks and zipping completed successfully.\n"
+        yield f"✅ {task_type.upper()} generated and zipped successfully", get_debug_info(), full_log, gr.update(value=ZIP_OUTPUT_PATH, visible=True)
+        
     except Exception as e:
         error_log = full_log + f"\n[Global Exception] Exception occurred:\n{str(e)}"
-        return "❌ Global exception during final generation", get_debug_info(), error_log, None
+        yield "❌ Global exception during final generation", get_debug_info(), error_log, gr.update(visible=False)
+
 
 # ==========================================
 # --- 🚀 UI Configuration & Advanced CSS ---
@@ -566,17 +604,18 @@ with gr.Blocks(theme=purple_theme, css=custom_css) as demo:
     
     parse_btn.click(fn=run_mineru_parsing_and_dag_gen, outputs=[parse_status, debug_view, cmd_logs])
 
-    def trigger_gen(task):
-        status, debug, logs, file_path = run_final_generation(task)
-        file_update = gr.update(value=file_path, visible=True) if file_path else gr.update(visible=False)
-        return status, debug, logs, file_update
+    # 为了确保 Gradio 生成器模式正确工作，拆分为显式的封装函数
+    def trigger_gen_ppt(): yield from run_final_generation("ppt")
+    def trigger_gen_poster(): yield from run_final_generation("poster")
+    def trigger_gen_pr(): yield from run_final_generation("pr")
+    def trigger_gen_all(): yield from run_final_generation("all")
 
-    gen_ppt_btn.click(fn=lambda: trigger_gen("ppt"), outputs=[gen_status, debug_view, cmd_logs, download_file])
-    gen_poster_btn.click(fn=lambda: trigger_gen("poster"), outputs=[gen_status, debug_view, cmd_logs, download_file])
-    gen_pr_btn.click(fn=lambda: trigger_gen("pr"), outputs=[gen_status, debug_view, cmd_logs, download_file])
-    gen_all_btn.click(fn=lambda: trigger_gen("all"), outputs=[gen_status, debug_view, cmd_logs, download_file])
+    gen_ppt_btn.click(fn=trigger_gen_ppt, outputs=[gen_status, debug_view, cmd_logs, download_file])
+    gen_poster_btn.click(fn=trigger_gen_poster, outputs=[gen_status, debug_view, cmd_logs, download_file])
+    gen_pr_btn.click(fn=trigger_gen_pr, outputs=[gen_status, debug_view, cmd_logs, download_file])
+    gen_all_btn.click(fn=trigger_gen_all, outputs=[gen_status, debug_view, cmd_logs, download_file])
     
     refresh_btn.click(fn=get_debug_info, outputs=debug_view)
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch()