speed up

main.py

@@ -99,7 +99,7 @@ def generate_paper_poster(
             # with open(output.replace(".json", ".html"), "w") as f:
             #     f.write(html)
             # take_screenshot(output, html)
-
+            print("海报生成成功!")
             return poster, html
 
         except Exception as e:

poster/poster.py

@@ -395,30 +395,61 @@ def generate_html_v2(vendor: str, model: str, poster: BaseModel, figures: list[s
     min_proportion = float('inf')
     min_html = None
     min_html_with_figures = None
-
-    while True:
-        body = re.search(r"```html\n(.*?)\n```", output, re.DOTALL).group(1)
-
+    min_body = None
+    min_poster_sizes = None
+    min_section_sizes = None
+
+    def generate_single_html(prompt_input):
+        """单个HTML生成函数，用于多线程执行"""
+        result_output = layout_chain.invoke(prompt_input).content
+        body = re.search(r"```html\n(.*?)\n```", result_output, re.DOTALL).group(1)
+        
         html = HTML_TEMPLATE.format(style=style, body=body)
         html_with_figures = replace_figures_in_html(html, figures)
-
+        
         poster_sizes = get_sizes("poster", html_with_figures)
         section_sizes = get_sizes("section", html_with_figures)
-
+        
         proportion = calculate_blank_proportion(poster_sizes, section_sizes)
         
-        print(f"当前比例: {proportion:.0%}")
+        return {
+            "body": body,
+            "html": html,
+            "html_with_figures": html_with_figures,
+            "poster_sizes": poster_sizes,
+            "section_sizes": section_sizes,
+            "proportion": proportion
+        }
+
+    # 初始生成两个HTML布局
+    prompt_inputs = [
+        {"style": style, "poster": poster},
+        {"style": style, "poster": poster}
+    ]
+    
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        initial_results = list(executor.map(generate_single_html, prompt_inputs))
+    
+    # 检查初始生成的两个结果
+    for result in initial_results:
+        proportion = result["proportion"]
+        print(f"初始生成比例: {proportion:.0%}")
+        
+        # 更新最佳结果
         if proportion < min_proportion:
             min_proportion = proportion
-            min_html = html
-            min_html_with_figures = html_with_figures
+            min_html = result["html"]
+            min_html_with_figures = result["html_with_figures"]
+            min_body = result["body"]
+            min_poster_sizes = result["poster_sizes"]
+            min_section_sizes = result["section_sizes"]
         
+        # 如果找到满足条件的结果，直接返回
         if proportion <= 0.1:
-            print(
-                f"Attempted {attempt} times, remaining {proportion:.0%} blank spaces."
-            )
-            return {"html": html, "html_with_figures": html_with_figures}
+            print(f"Initial generation successful, remaining {proportion:.0%} blank spaces.")
+            return {"html": result["html"], "html_with_figures": result["html_with_figures"]}
 
+    while True:
         attempt += 1
         if attempt > max_attempts:
             if min_proportion <= 0.2:
@@ -429,26 +460,51 @@ def generate_html_v2(vendor: str, model: str, poster: BaseModel, figures: list[s
             else:
                 raise ValueError(f"Invalid blank spaces: {min_proportion:.0%}")
 
-
+        # 基于最好的结果生成两个新的
         react = [
             HumanMessage(
                 content=f"""# Previous Body
-{body}
+{min_body}
 
 # Previous Size of Columns in Poster
-{poster_sizes}
+{min_poster_sizes}
 
 # Previous Size of Columns in Section
-{section_sizes}
+{min_section_sizes}
 
-Now there are {proportion:.0%} blank spaces. Please regenerate the content to create a more balanced poster layout.
+Now there are {min_proportion:.0%} blank spaces. Please regenerate the content to create a more balanced poster layout.
 """
             ),
         ]
-
-        output = layout_chain.invoke(
+        prompt_inputs = [
+            {"style": style, "poster": poster, "react": react},
             {"style": style, "poster": poster, "react": react}
-        ).content
+        ]
+
+        # 使用多线程同时生成两个HTML
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            results = list(executor.map(generate_single_html, prompt_inputs))
+
+        # 检查两个结果
+        for result in results:
+            proportion = result["proportion"]
+            print(f"当前比例: {proportion:.0%}")
+            
+            # 更新最佳结果
+            if proportion < min_proportion:
+                min_proportion = proportion
+                min_html = result["html"]
+                min_html_with_figures = result["html_with_figures"]
+                min_body = result["body"]
+                min_poster_sizes = result["poster_sizes"]
+                min_section_sizes = result["section_sizes"]
+            
+            # 如果找到满足条件的结果，直接返回
+            if proportion <= 0.1:
+                print(
+                    f"Attempted {attempt} times, remaining {proportion:.0%} blank spaces."
+                )
+                return {"html": result["html"], "html_with_figures": result["html_with_figures"]}
 
 
 # def take_screenshot(output: str, html: str):