Update UI to support up to 4 calligrapher/style pairs

app.py

@@ -273,13 +273,14 @@ def parse_font_style(font_style: str) -> str:
 # Keep it lazy inside the @spaces.GPU worker to avoid any pre-fork CUDA side effects.
 
 
-def _get_generation_duration(text, font, author, num_steps, start_seed, num_images, progress=None):
+def _get_generation_duration(text, pairs, num_steps, start_seed, num_images, progress=None):
     """Calculate dynamic GPU duration: 24s base + 3s per image"""
-    return 24 + int(3 * num_images)
+    num_pairs = len(pairs) if pairs else 1
+    return 24 + int(3 * num_images * num_pairs)
 
 
 @spaces.GPU(duration=_get_generation_duration)
-def run_generation(text, font, author, num_steps, start_seed, num_images, progress=gr.Progress()):
+def run_generation(text, pairs, num_steps, start_seed, num_images, progress=gr.Progress()):
     """
     Load model, apply FP8 quantization, and generate images.
     All in one GPU session to avoid redundant loading.
@@ -328,21 +329,32 @@ def run_generation(text, font, author, num_steps, start_seed, num_images, progre
     logger.info("Model weights decompressed to bfloat16 upon load. Skipping dynamic quantization to ensure stability.")
     
     # Step 3: Generate images
-    logger.info(f"Generating {num_images} images...")
+    total_gens = len(pairs) * num_images
+    logger.info(f"Generating {total_gens} images across {len(pairs)} styles...")
     results = []
     seeds_used = []
-    for i in range(num_images):
-        loop_progress = 0.82 + ((i + 1) / max(num_images, 1)) * 0.16
-        progress(loop_progress, desc=f"生成第 {i+1}/{num_images} 张 / Generating {i+1}/{num_images}")
-        current_seed = start_seed + i
-        result_img, cond_img = gen.generate(
-            text=text, font_style=font, author=author,
-            num_steps=num_steps, seed=current_seed,
-            guidance=1.0,
-        )
-        results.append((result_img, f"Seed: {current_seed}"))
-        seeds_used.append(current_seed)
-        logger.info(f"  Generated image {i+1}/{num_images}")
+    
+    gen_idx = 0
+    for author, font in pairs:
+        for i in range(num_images):
+            gen_idx += 1
+            loop_progress = 0.82 + (gen_idx / max(total_gens, 1)) * 0.16
+            progress(loop_progress, desc=f"生成第 {gen_idx}/{total_gens} 张 / Generating {gen_idx}/{total_gens}")
+            current_seed = start_seed + i
+            
+            cond_author = author if author != "None (Synthetic / 合成风格)" else None
+            
+            result_img, cond_img = gen.generate(
+                text=text, font_style=font, author=cond_author,
+                num_steps=num_steps, seed=current_seed,
+                guidance=1.0,
+            )
+            
+            author_label = author if author else "Synthetic"
+            label = f"{author_label} - {font} (Seed: {current_seed})"
+            results.append((result_img, label))
+            seeds_used.append(current_seed)
+            logger.info(f"  Generated image {gen_idx}/{total_gens}")
     
     progress(1.0, desc="生成完成 / Generation complete")
     return results, seeds_used
@@ -350,8 +362,7 @@ def run_generation(text, font, author, num_steps, start_seed, num_images, progre
 
 def interactive_session(
     text: str,
-    author_dropdown: str,
-    font_style: str,
+    a1, f1, a2, f2, a3, f3, a4, f4,
     num_steps: int,
     start_seed: int,
     num_images: int,
@@ -366,13 +377,16 @@ def interactive_session(
     if len(text) > 7:
         raise gr.Error(f"文本最多7个字符 / Text must be at most 7 characters. Current: {len(text)}")
     
-    # Parse font style
-    font = parse_font_style(font_style)
-    if font is None:
-        raise gr.Error(f"无法识别的字体风格 / Unknown font style: {font_style}")
-    
-    # Determine author
-    author = author_dropdown if author_dropdown != "None (Synthetic / 合成风格)" else None
+    raw_pairs = [(a1, f1), (a2, f2), (a3, f3), (a4, f4)]
+    pairs = []
+    for a, f_style in raw_pairs:
+        if a and f_style:
+            parsed_font = parse_font_style(f_style)
+            if parsed_font is not None:
+                pairs.append((a, parsed_font))
+                
+    if not pairs:
+        raise gr.Error("请至少选择一项书法家和字体组合 / Please select at least one combination")
     
     # Run generation (includes model loading + FP8 quantization + generation)
     yield "⏳ 队列中：准备任务... / Queued: preparing task...", []
@@ -387,14 +401,15 @@ def interactive_session(
     progress(0.22, desc="进入生成阶段 / Entering generation stage...")
 
     results, seeds_used = run_generation(
-        text, font, author, num_steps, start_seed, num_images, progress
+        text, pairs, num_steps, start_seed, num_images, progress
     )
     
     progress(1.0, desc="完成!")
     
     # Final status
-    if num_images > 1:
-        final_status = f"✅ 全部完成！共 {num_images} 张 (Seeds: {seeds_used[0]}-{seeds_used[-1]})"
+    total_imgs = len(results)
+    if total_imgs > 1:
+        final_status = f"✅ 全部完成！共 {total_imgs} 张 (Seeds: {seeds_used[0]}-{seeds_used[-1]})"
     else:
         final_status = f"✅ 完成！Seed: {seeds_used[0]}"
     yield final_status, results
@@ -428,31 +443,35 @@ with gr.Blocks(title="UniCalli - Chinese Calligraphy Generator / 中国书法生
                 max_lines=1
             )
             
-            gr.Markdown("### 👤 书法家选择 / Calligrapher Selection")
+            gr.Markdown("### 👤 书法家与字体组合 / Calligraphers & Fonts (最多4组 / Up to 4)")
             
-            author_dropdown = gr.Dropdown(
-                label="1. 选择书法家 / Select Calligrapher",
-                choices=["None (Synthetic / 合成风格)"] + AUTHOR_LIST,
-                value="文征明",
-                info="先选择历史书法家 / Choose a historical calligrapher first"
-            )
+            author_dropdowns = []
+            font_dropdowns = []
             
-            # Get initial fonts for default author (文征明)
             initial_author = "文征明"
             initial_fonts = AUTHOR_FONTS.get(initial_author, ["楷", "草", "行"])
             initial_font_choices = [FONT_STYLE_NAMES[f] for f in initial_fonts if f in FONT_STYLE_NAMES]
-            # Default to first available font for the author, prefer "行" for 文征明
             if initial_author == "文征明" and "行" in initial_fonts:
                 default_font = FONT_STYLE_NAMES["行"]
             else:
                 default_font = initial_font_choices[0] if initial_font_choices else "草 (Cursive Script)"
             
-            font_style = gr.Dropdown(
-                label="2. 选择字体风格 / Select Font Style",
-                choices=initial_font_choices,
-                value=default_font,
-                info="根据所选书法家显示可用字体 / Shows available fonts for selected calligrapher"
-            )
+            for i in range(4):
+                with gr.Group():
+                    gr.Markdown(f"**组合 {i+1} / Combination {i+1}** (不填则忽略 / Leave blank to ignore)")
+                    with gr.Row():
+                        a_drop = gr.Dropdown(
+                            label=f"书法家 / Calligrapher",
+                            choices=["None (Synthetic / 合成风格)"] + AUTHOR_LIST,
+                            value="文征明" if i == 0 else None,
+                        )
+                        f_drop = gr.Dropdown(
+                            label=f"字体风格 / Font Style",
+                            choices=initial_font_choices if i == 0 else list(FONT_STYLE_NAMES.values()),
+                            value=default_font if i == 0 else None,
+                        )
+                        author_dropdowns.append(a_drop)
+                        font_dropdowns.append(f_drop)
             
             gr.Markdown("### ⚙️ 生成设置 / Generation Settings")
             
@@ -524,23 +543,23 @@ with gr.Blocks(title="UniCalli - Chinese Calligraphy Generator / 中国书法生
         gr.Markdown(author_info_md)
     
     # Event handlers
-    author_dropdown.change(
-        fn=update_font_choices,
-        inputs=[author_dropdown],
-        outputs=[font_style]
-    )
+    for i in range(4):
+        author_dropdowns[i].change(
+            fn=update_font_choices,
+            inputs=[author_dropdowns[i]],
+            outputs=[font_dropdowns[i]]
+        )
+    
+    # Prepare inputs list for the interactive session
+    session_inputs = [text_input]
+    for i in range(4):
+        session_inputs.extend([author_dropdowns[i], font_dropdowns[i]])
+    session_inputs.extend([num_steps, start_seed, num_images])
     
     # Generate button - uses streaming for live updates
     generate_btn.click(
         fn=interactive_session,
-        inputs=[
-            text_input,
-            author_dropdown,
-            font_style,
-            num_steps,
-            start_seed,
-            num_images,
-        ],
+        inputs=session_inputs,
         outputs=[status_text, output_gallery]
     )
     
@@ -548,20 +567,11 @@ with gr.Blocks(title="UniCalli - Chinese Calligraphy Generator / 中国书法生
     gr.Markdown("### 📋 示例 / Examples")
     gr.Examples(
         examples=[
-            ["相见时难别亦难", "文征明", "行 (Running Script)", 4, 1024, 1],
-            ["天道酬勤", "王羲之", "草 (Cursive Script)", 4, 42, 1],
-            ["厚德载物", "赵孟頫", "楷 (Regular Script)", 4, 123, 1],
-            ["海内存知己", "黄庭坚", "行 (Running Script)", 4, 456, 1],
-            ["宁静致远", "None (Synthetic / 合成风格)", "楷 (Regular Script)", 4, 789, 1],
-        ],
-        inputs=[
-            text_input,
-            author_dropdown,
-            font_style,
-            num_steps,
-            start_seed,
-            num_images,
+            ["相见时难别亦难", "文征明", "行 (Running Script)", None, None, None, None, None, None, 4, 1024, 1],
+            ["天道酬勤", "王羲之", "草 (Cursive Script)", "黄庭坚", "行 (Running Script)", None, None, None, None, 4, 42, 1],
+            ["厚德载物", "赵孟頫", "楷 (Regular Script)", None, None, None, None, None, None, 4, 123, 1],
         ],
+        inputs=session_inputs,
     )