From ae5a3b8f1d4f3a479783757e0cc915b7ba8df057 Mon Sep 17 00:00:00 2001
From: xsl <youmiss@163.com>
Date: Sun, 26 Jul 2026 16:12:43 +0800
Subject: [PATCH] =?UTF-8?q?refactor(comfyui):=20=E6=A8=A1=E5=9E=8B?=
 =?UTF-8?q?=E5=88=87=E6=8D=A2=E6=97=B6=E5=90=8C=E6=AD=A5VAE=20+=20?=
 =?UTF-8?q?=E5=AE=8C=E5=96=84=E7=BC=96=E7=A0=81=E5=99=A8=E6=98=A0=E5=B0=84?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 新增 _vae_for_unet() 按模型自动切换 VAE(Z-Image用ae, Flux.2用flux2-vae)
- _clip_for_unet 增加 z-image 识别
- 为将来多模型切换做准备(当前 Q4 不受影响)
---
 hairline/comfyui.py | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/hairline/comfyui.py b/hairline/comfyui.py
index ee42ddf..1e5528f 100644
--- a/hairline/comfyui.py
+++ b/hairline/comfyui.py
@@ -32,18 +32,28 @@ _SEED_NODE = "6"       # RandomNoise
 _PROMPT_NODE = "60"    # JjkText：提示词
 _UNET_NODE = "16"      # UNETLoader / UnetLoaderGGUF：Flux 模型加载
 _CLIP_NODE = "61"      # CLIPLoader：qwen 文本编码器
+_VAE_NODE = "3"        # VAELoader
 
 # Flux 模型 → 配套文本编码器映射。切换 unet 时自动同步编码器，避免维度不匹配。
-# 规则：4b 系列配 qwen_3_4b，9b 系列（fp8/GGUF）配 qwen_3_8b_fp8mixed。
 def _clip_for_unet(unet_name: str) -> str | None:
     """根据 unet 文件名推断配套的文本编码器文件名；无法推断返回 None。"""
     low = unet_name.lower()
-    if "4b" in low and "9b" not in low:
-        return "qwen_3_4b.safetensors"
-    if "9b" in low:
+    if "9b" in low:                       # Flux.2 9B 系列
         return "qwen_3_8b_fp8mixed.safetensors"
+    if "z-image" in low:                  # Z-Image-Turbo 用 4B 编码器
+        return "qwen_3_4b.safetensors"
+    if "4b" in low:                       # Flux.2 4B
+        return "qwen_3_4b.safetensors"
     return None
 
+# Flux 模型 → 配套 VAE 映射。Z-Image 用 ae.safetensors，Flux.2 系列用 flux2-vae。
+def _vae_for_unet(unet_name: str) -> str | None:
+    """根据 unet 文件名推断配套 VAE 文件名；无法推断返回 None（保持工作流原值）。"""
+    low = unet_name.lower()
+    if "z-image" in low:
+        return "ae.safetensors"
+    return None  # Flux.2 系列 vae 在工作流里已正确配置，不覆盖
+
 _wf_cache: dict[str, dict] = {}      # path → workflow JSON
 _wf_output_node: dict[str, str] = {}  # path → SaveImage 节点 ID
 
@@ -153,6 +163,11 @@ def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT, prompt: str = Non
             clip_name = _clip_for_unet(unet_name)
             if clip_node is not None and clip_name is not None:
                 clip_node["inputs"]["clip_name"] = clip_name
+            # 同步切换 VAE（Z-Image 用 ae.safetensors，Flux.2 保持 flux2-vae）
+            vae_node = wf.get(_VAE_NODE)
+            vae_name = _vae_for_unet(unet_name)
+            if vae_node is not None and vae_name is not None:
+                vae_node["inputs"]["vae_name"] = vae_name
 
         # 诊断：落盘实际提交的工作流 + 输入图，便于和手动 ComfyUI 跑的对比
         try: