refactor(comfyui): 模型切换时同步VAE + 完善编码器映射
- 新增 _vae_for_unet() 按模型自动切换 VAE(Z-Image用ae, Flux.2用flux2-vae) - _clip_for_unet 增加 z-image 识别 - 为将来多模型切换做准备(当前 Q4 不受影响)
This commit is contained in:
+19
-4
@@ -32,18 +32,28 @@ _SEED_NODE = "6" # RandomNoise
|
|||||||
_PROMPT_NODE = "60" # JjkText:提示词
|
_PROMPT_NODE = "60" # JjkText:提示词
|
||||||
_UNET_NODE = "16" # UNETLoader / UnetLoaderGGUF:Flux 模型加载
|
_UNET_NODE = "16" # UNETLoader / UnetLoaderGGUF:Flux 模型加载
|
||||||
_CLIP_NODE = "61" # CLIPLoader:qwen 文本编码器
|
_CLIP_NODE = "61" # CLIPLoader:qwen 文本编码器
|
||||||
|
_VAE_NODE = "3" # VAELoader
|
||||||
|
|
||||||
# Flux 模型 → 配套文本编码器映射。切换 unet 时自动同步编码器,避免维度不匹配。
|
# Flux 模型 → 配套文本编码器映射。切换 unet 时自动同步编码器,避免维度不匹配。
|
||||||
# 规则:4b 系列配 qwen_3_4b,9b 系列(fp8/GGUF)配 qwen_3_8b_fp8mixed。
|
|
||||||
def _clip_for_unet(unet_name: str) -> str | None:
|
def _clip_for_unet(unet_name: str) -> str | None:
|
||||||
"""根据 unet 文件名推断配套的文本编码器文件名;无法推断返回 None。"""
|
"""根据 unet 文件名推断配套的文本编码器文件名;无法推断返回 None。"""
|
||||||
low = unet_name.lower()
|
low = unet_name.lower()
|
||||||
if "4b" in low and "9b" not in low:
|
if "9b" in low: # Flux.2 9B 系列
|
||||||
return "qwen_3_4b.safetensors"
|
|
||||||
if "9b" in low:
|
|
||||||
return "qwen_3_8b_fp8mixed.safetensors"
|
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
|
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_cache: dict[str, dict] = {} # path → workflow JSON
|
||||||
_wf_output_node: dict[str, str] = {} # path → SaveImage 节点 ID
|
_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)
|
clip_name = _clip_for_unet(unet_name)
|
||||||
if clip_node is not None and clip_name is not None:
|
if clip_node is not None and clip_name is not None:
|
||||||
clip_node["inputs"]["clip_name"] = clip_name
|
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 跑的对比
|
# 诊断:落盘实际提交的工作流 + 输入图,便于和手动 ComfyUI 跑的对比
|
||||||
try:
|
try:
|
||||||
|
|||||||
Reference in New Issue
Block a user