This commit is contained in:
AI_Architect
2026-07-02 13:55:15 +08:00
parent 80b5e0873a
commit 6fea60ba50
6 changed files with 1098 additions and 1829 deletions
+52 -22
View File
@@ -71,11 +71,17 @@ try:
# 注册视频转码端点(预览窗口实时缩放依赖)
from . import xb_video_server
from .nodes_rocm import (XB_ROCmKSampler, XB_ROCmKSamplerAdvanced,
XB_ROCmSamplerCustom, XB_ROCmSamplerCustomAdvanced,
XB_ROCmVAEDecode, XB_ROCmVAEEncode,
XB_ROCmVAEDecodeTemporal, XB_ROCmLTXVAEDecode,
XB_ROCmMemCleaner)
from .nodes_rocm import XB_ROCmMemCleaner
from .nodes_vanilla_wrappers import (XB_KSampler, XB_KSamplerAdvanced,
XB_SamplerCustom, XB_SamplerCustomAdvanced,
XB_VAEDecode, XB_VAEDecodeTiled,
XB_VAEDecodeTiledImage,
XB_VAEEncode, XB_VAEEncodeTiled,
XB_VAEEncodeForInpaint,
_AliasKSampler, _AliasKSamplerAdvanced,
_AliasSamplerCustom, _AliasSamplerCustomAdvanced,
_AliasVAEDecode, _AliasVAEDecodeTemporal,
_AliasLTXVAEDecode, _AliasVAEEncode)
from .nodes_string_merge import XB_StringMerge
from .nodes_msr import XB_MSR
from .nodes_comic import XB_ComicPromptParser, XB_ComicTextRenderer, XB_AutoBubbleTextRenderer
@@ -121,15 +127,27 @@ try:
"XB_Video_Merger": XB_Video_Merger,
"XB_StoryboardSlicer": XB_StoryboardSlicer,
"XB_SageAttentionAccelerator": XB_SageAttentionAccelerator,
"XB_ROCmKSampler": XB_ROCmKSampler,
"XB_ROCmKSamplerAdvanced": XB_ROCmKSamplerAdvanced,
"XB_ROCmSamplerCustom": XB_ROCmSamplerCustom,
"XB_ROCmSamplerCustomAdvanced": XB_ROCmSamplerCustomAdvanced,
"XB_ROCmVAEDecode": XB_ROCmVAEDecode,
"XB_ROCmVAEEncode": XB_ROCmVAEEncode,
"XB_ROCmVAEDecodeTemporal": XB_ROCmVAEDecodeTemporal,
"XB_ROCmLTXVAEDecode": XB_ROCmLTXVAEDecode,
"XB_ROCmMemCleaner": XB_ROCmMemCleaner,
# ── 借尸还魂:旧 ROCm 节点 ID → 别名类(兼容旧参数名) ──
"XB_ROCmKSampler": _AliasKSampler,
"XB_ROCmKSamplerAdvanced": _AliasKSamplerAdvanced,
"XB_ROCmSamplerCustom": _AliasSamplerCustom,
"XB_ROCmSamplerCustomAdvanced": _AliasSamplerCustomAdvanced,
"XB_ROCmVAEDecode": _AliasVAEDecode,
"XB_ROCmVAEEncode": _AliasVAEEncode,
"XB_ROCmVAEDecodeTemporal": _AliasVAEDecodeTemporal,
"XB_ROCmLTXVAEDecode": _AliasLTXVAEDecode,
# ── 新原版优化节点 ──
"XB_KSampler": XB_KSampler,
"XB_KSamplerAdvanced": XB_KSamplerAdvanced,
"XB_SamplerCustom": XB_SamplerCustom,
"XB_SamplerCustomAdvanced": XB_SamplerCustomAdvanced,
"XB_VAEDecode": XB_VAEDecode,
"XB_VAEDecodeTiled": XB_VAEDecodeTiled,
"XB_VAEDecodeTiledImage": XB_VAEDecodeTiledImage,
"XB_VAEEncode": XB_VAEEncode,
"XB_VAEEncodeTiled": XB_VAEEncodeTiled,
"XB_VAEEncodeForInpaint": XB_VAEEncodeForInpaint,
"XB_WanT5Loader": XB_WanT5Loader,
"XB_WanCompileSettings": XB_WanCompileSettings,
"XB_WanModelLoader": XB_WanModelLoader,
@@ -187,15 +205,27 @@ try:
"XB_Video_Merger": "XB-BOX - Video Merger",
"XB_StoryboardSlicer": "XB-BOX - Storyboard Slicer",
"XB_SageAttentionAccelerator": "XB-BOX - SageAttention Accelerator",
"XB_ROCmKSampler": "XB-BOX - 🚀 ROCm 采样器",
"XB_ROCmKSamplerAdvanced": "XB-BOX - 🚀 ROCm 高级采样器",
"XB_ROCmSamplerCustom": "XB-BOX - ⚙️ ROCm 自定义采样器",
"XB_ROCmSamplerCustomAdvanced": "XB-BOX - ⚙️ ROCm 自定义高级采样器",
"XB_ROCmVAEDecode": "XB-BOX - 🖼️ ROCm VAE 解码",
"XB_ROCmVAEEncode": "XB-BOX - 📦 ROCm VAE 编码",
"XB_ROCmVAEDecodeTemporal": "XB-BOX - 🎬 ROCm VAE 时空解码",
"XB_ROCmLTXVAEDecode": "XB-BOX - 🎬 ROCm LTX VAE 时空解码",
"XB_ROCmMemCleaner": "XB-BOX - 🧹 ROCm 显存清理",
"XB_ROCmMemCleaner": "XB-BOX - 🧹 显存清理",
# ── 借尸还魂:旧节点名 + 新的优化版显示名 ──
"XB_ROCmKSampler": "XB-BOX - 采样器(原版优化)",
"XB_ROCmKSamplerAdvanced": "XB-BOX - 高级采样器(原版优化)",
"XB_ROCmSamplerCustom": "XB-BOX - 自定义采样器(原版优化)",
"XB_ROCmSamplerCustomAdvanced": "XB-BOX - 自定义高级采样器(原版优化)",
"XB_ROCmVAEDecode": "XB-BOX - VAE解码(原版优化)",
"XB_ROCmVAEEncode": "XB-BOX - VAE编码(原版优化)",
"XB_ROCmVAEDecodeTemporal": "XB-BOX - VAE分块解码(原版优化)",
"XB_ROCmLTXVAEDecode": "XB-BOX - VAE分块解码(原版优化)",
# ── 新优化版节点 ──
"XB_KSampler": "XB-BOX - 采样器(原版优化)",
"XB_KSamplerAdvanced": "XB-BOX - 高级采样器(原版优化)",
"XB_SamplerCustom": "XB-BOX - 自定义采样器(原版优化)",
"XB_SamplerCustomAdvanced": "XB-BOX - 自定义高级采样器(原版优化)",
"XB_VAEDecode": "XB-BOX - VAE解码(原版优化)",
"XB_VAEDecodeTiled": "XB-BOX - VAE分块解码(原版优化)",
"XB_VAEDecodeTiledImage": "XB-BOX - VAE解码(原版优化)",
"XB_VAEEncode": "XB-BOX - VAE编码(原版优化)",
"XB_VAEEncodeTiled": "XB-BOX - VAE分块编码(原版优化)",
"XB_VAEEncodeForInpaint": "XB-BOX - VAE修补编码(原版优化)",
"XB_WanT5Loader": "XB-BOX - 📝 Wan T5 加载器(FP8)",
"XB_WanCompileSettings": "XB-BOX - ⚡ Wan 编译设置",
"XB_WanModelLoader": "XB-BOX - 🧠 Wan 模型加载",
+215 -259
View File
@@ -1087,297 +1087,122 @@
}
},
"XB_ROCmKSampler": {
"display_name": "XB-BOX - 🚀 ROCm 采样器",
"display_name": "XB-BOX - 采样器(原版优化)",
"inputs": {
"model": {
"name": "模型"
},
"seed": {
"name": "随机种子"
},
"steps": {
"name": "采样步数"
},
"cfg": {
"name": "引导系数(CFG)"
},
"sampler": {
"name": "采样器"
},
"scheduler": {
"name": "调度器"
},
"positive": {
"name": "正面提示词"
},
"negative": {
"name": "负面提示词"
},
"latent": {
"name": "潜变量"
},
"denoise": {
"name": "降噪强度"
},
"cleanup": {
"name": "采样后清理"
}
"model": { "name": "模型" },
"seed": { "name": "种子" },
"steps": { "name": "步数" },
"cfg": { "name": "CFG" },
"sampler": { "name": "采样器" },
"scheduler": { "name": "调度器" },
"positive": { "name": "正面条件" },
"negative": { "name": "负面条件" },
"latent": { "name": "潜空间" },
"denoise": { "name": "降噪" },
"cleanup": { "name": "执行前清理" }
},
"outputs": {
"0": {
"name": "潜变量"
}
}
"outputs": { "0": { "name": "潜空间" } }
},
"XB_ROCmKSamplerAdvanced": {
"display_name": "XB-BOX - 🚀 ROCm 高级采样器",
"display_name": "XB-BOX - 高级采样器(原版优化)",
"inputs": {
"model": {
"name": "模型"
},
"add_noise": {
"name": "添加噪声"
},
"noise_seed": {
"name": "噪声种子"
},
"steps": {
"name": "采样步数"
},
"cfg": {
"name": "引导系数(CFG)"
},
"sampler": {
"name": "采样器"
},
"scheduler": {
"name": "调度器"
},
"positive": {
"name": "正面提示词"
},
"negative": {
"name": "负面提示词"
},
"latent": {
"name": "潜变量"
},
"start_at_step": {
"name": "起始步数"
},
"end_at_step": {
"name": "结束步数"
},
"return_with_leftover_noise": {
"name": "返回残留噪声"
},
"cleanup": {
"name": "采样后清理"
}
"model": { "name": "模型" },
"add_noise": { "name": "加噪" },
"noise_seed": { "name": "噪声种子" },
"steps": { "name": "步数" },
"cfg": { "name": "CFG" },
"sampler": { "name": "采样器" },
"scheduler": { "name": "调度器" },
"positive": { "name": "正面条件" },
"negative": { "name": "负面条件" },
"latent": { "name": "潜空间" },
"start_at_step": { "name": "起始步" },
"end_at_step": { "name": "结束步" },
"return_with_leftover_noise": { "name": "残留噪声" },
"cleanup": { "name": "执行前清理" }
},
"outputs": {
"0": {
"name": "潜变量"
}
}
"outputs": { "0": { "name": "潜空间" } }
},
"XB_ROCmSamplerCustom": {
"display_name": "XB-BOX - ⚙️ ROCm 自定义采样器",
"display_name": "XB-BOX - 自定义采样器(原版优化)",
"inputs": {
"model": {
"name": "模型"
},
"add_noise": {
"name": "添加噪声"
},
"noise_seed": {
"name": "噪声种子"
},
"cfg": {
"name": "引导系数(CFG)"
},
"positive": {
"name": "正面提示词"
},
"negative": {
"name": "负面提示词"
},
"sampler": {
"name": "采样器"
},
"sigmas": {
"name": "西格玛"
},
"latent_image": {
"name": "Latent图像"
},
"cleanup": {
"name": "采样后清理"
}
"model": { "name": "模型" },
"add_noise": { "name": "加噪" },
"noise_seed": { "name": "噪声种子" },
"cfg": { "name": "CFG" },
"positive": { "name": "正面条件" },
"negative": { "name": "负面条件" },
"sampler": { "name": "采样器" },
"sigmas": { "name": "Sigma" },
"latent_image": { "name": "潜空间" },
"cleanup": { "name": "执行前清理" }
},
"outputs": {
"0": {
"name": "输出Latent"
},
"1": {
"name": "降噪Latent"
}
}
"outputs": { "0": { "name": "输出" }, "1": { "name": "降噪输出" } }
},
"XB_ROCmSamplerCustomAdvanced": {
"display_name": "XB-BOX - ⚙️ ROCm 自定义高级采样器",
"display_name": "XB-BOX - 自定义高级采样器(原版优化)",
"inputs": {
"noise": {
"name": "噪波"
},
"guider": {
"name": "引导器"
},
"sampler": {
"name": "采样器"
},
"sigmas": {
"name": "西格玛"
},
"latent_image": {
"name": "Latent图像"
},
"cleanup": {
"name": "采样后清理"
}
"noise": { "name": "噪声" },
"guider": { "name": "引导器" },
"sampler": { "name": "采样器" },
"sigmas": { "name": "Sigma" },
"latent_image": { "name": "潜空间" },
"cleanup": { "name": "执行前清理" }
},
"outputs": {
"0": {
"name": "输出Latent"
},
"1": {
"name": "降噪Latent"
}
}
"outputs": { "0": { "name": "输出" }, "1": { "name": "降噪输出" } }
},
"XB_ROCmVAEDecode": {
"display_name": "XB-BOX - 🖼️ ROCm VAE 解码",
"display_name": "XB-BOX - VAE解码(原版优化)",
"inputs": {
"samples": {
"name": "潜变量"
},
"vae": {
"name": "VAE模型"
},
"tile": {
"name": "空间分块(0=自动)"
},
"overlap": {
"name": "分块重叠(0=自动)"
},
"cleanup": {
"name": "清理力度"
}
"samples": { "name": "潜空间" },
"vae": { "name": "VAE" },
"tile": { "name": "分块大小(0=自动)" },
"overlap": { "name": "重叠(0=自动)" },
"cleanup": { "name": "执行前清理" }
},
"outputs": {
"0": {
"name": "图像"
}
}
"outputs": { "0": { "name": "图像" } }
},
"XB_ROCmVAEEncode": {
"display_name": "XB-BOX - 📦 ROCm VAE 编码",
"display_name": "XB-BOX - VAE编码(原版优化)",
"inputs": {
"pixels": {
"name": "图像"
},
"vae": {
"name": "VAE模型"
},
"tile": {
"name": "空间分块(0=自动)"
},
"overlap": {
"name": "分块重叠(0=自动)"
},
"temporal_size": {
"name": "时间分块帧数(0=自动)"
},
"temporal_overlap": {
"name": "时间重叠(0=自动)"
},
"cleanup": {
"name": "编码后清理"
}
"pixels": { "name": "图像" },
"vae": { "name": "VAE" },
"tile": { "name": "分块大小(0=自动)" },
"overlap": { "name": "重叠(0=自动)" },
"temporal_size": { "name": "时间分块(0=自动)" },
"temporal_overlap": { "name": "时间重叠(0=自动)" },
"cleanup": { "name": "执行前清理" }
},
"outputs": {
"0": {
"name": "潜变量"
}
}
"outputs": { "0": { "name": "潜空间" } }
},
"XB_ROCmVAEDecodeTemporal": {
"display_name": "XB-BOX - 🎬 ROCm VAE 时空解码",
"display_name": "XB-BOX - VAE分块解码(原版优化)",
"inputs": {
"samples": {
"name": "潜变量"
},
"vae": {
"name": "VAE模型"
},
"tile": {
"name": "空间分块(0=自动)"
},
"overlap": {
"name": "空间重叠(0=自动)"
},
"t_tile": {
"name": "时间分块帧数(0=自动)"
},
"t_overlap": {
"name": "时间重叠(0=自动)"
},
"cleanup": {
"name": "清理力度"
}
"samples": { "name": "潜空间" },
"vae": { "name": "VAE" },
"tile": { "name": "空间分块(0=自动)" },
"overlap": { "name": "空间重叠(0=自动)" },
"t_tile": { "name": "时间分块(0=自动)" },
"t_overlap": { "name": "时间重叠(0=自动)" },
"cleanup": { "name": "执行前清理" }
},
"outputs": {
"0": {
"name": "图像"
}
}
"outputs": { "0": { "name": "图像" } }
},
"XB_ROCmLTXVAEDecode": {
"display_name": "XB-BOX - 🎬 ROCm LTX VAE 时空解码",
"display_name": "XB-BOX - VAE分块解码(原版优化)",
"inputs": {
"samples": {
"name": "潜变量"
},
"vae": {
"name": "LTX VAE模型"
},
"spatial_tiles": {
"name": "空间分块数(1=不分块)"
},
"spatial_overlap": {
"name": "空间重叠(0=自动)"
},
"temporal_tile_length": {
"name": "时间分块长度(latent帧)"
},
"temporal_overlap": {
"name": "时间重叠(latent帧)"
},
"last_frame_fix": {
"name": "末帧修复"
},
"cleanup": {
"name": "清理力度"
}
"samples": { "name": "潜空间" },
"vae": { "name": "VAE" },
"tile": { "name": "空间分块(0=自动)" },
"overlap": { "name": "空间重叠(0=自动)" },
"t_tile": { "name": "时间分块(0=自动)" },
"t_overlap": { "name": "时间重叠(0=自动)" },
"cleanup": { "name": "执行前清理" }
},
"outputs": {
"0": {
"name": "图像"
}
}
"outputs": { "0": { "name": "图像" } }
},
"XB_ROCmMemCleaner": {
"display_name": "XB-BOX - 🧹 ROCm 显存清理",
"display_name": "XB-BOX - 显存清理",
"inputs": {
"mode": {
"name": "清理模式"
@@ -3289,5 +3114,136 @@
"name": "检测信息"
}
}
},
"XB_KSampler": {
"display_name": "XB-BOX - 采样器(原版优化)",
"inputs": {
"model": { "name": "模型" },
"seed": { "name": "种子" },
"steps": { "name": "步数" },
"cfg": { "name": "CFG" },
"sampler_name": { "name": "采样器" },
"scheduler": { "name": "调度器" },
"positive": { "name": "正面条件" },
"negative": { "name": "负面条件" },
"latent_image": { "name": "潜空间" },
"denoise": { "name": "降噪" },
"cleanup": { "name": "执行前清理" }
},
"outputs": { "0": { "name": "潜空间" } }
},
"XB_KSamplerAdvanced": {
"display_name": "XB-BOX - 高级采样器(原版优化)",
"inputs": {
"model": { "name": "模型" },
"add_noise": { "name": "加噪" },
"noise_seed": { "name": "噪声种子" },
"steps": { "name": "步数" },
"cfg": { "name": "CFG" },
"sampler_name": { "name": "采样器" },
"scheduler": { "name": "调度器" },
"positive": { "name": "正面条件" },
"negative": { "name": "负面条件" },
"latent_image": { "name": "潜空间" },
"start_at_step": { "name": "起始步" },
"end_at_step": { "name": "结束步" },
"return_with_leftover_noise": { "name": "返回残留噪声" },
"cleanup": { "name": "执行前清理" }
},
"outputs": { "0": { "name": "潜空间" } }
},
"XB_SamplerCustom": {
"display_name": "XB-BOX - 自定义采样器(原版优化)",
"inputs": {
"model": { "name": "模型" },
"add_noise": { "name": "加噪" },
"noise_seed": { "name": "噪声种子" },
"cfg": { "name": "CFG" },
"positive": { "name": "正面条件" },
"negative": { "name": "负面条件" },
"sampler": { "name": "采样器" },
"sigmas": { "name": "Sigma" },
"latent_image": { "name": "潜空间" },
"cleanup": { "name": "执行前清理" }
},
"outputs": { "0": { "name": "输出" }, "1": { "name": "降噪输出" } }
},
"XB_SamplerCustomAdvanced": {
"display_name": "XB-BOX - 自定义高级采样器(原版优化)",
"inputs": {
"noise": { "name": "噪声" },
"guider": { "name": "引导器" },
"sampler": { "name": "采样器" },
"sigmas": { "name": "Sigma" },
"latent_image": { "name": "潜空间" },
"cleanup": { "name": "执行前清理" }
},
"outputs": { "0": { "name": "输出" }, "1": { "name": "降噪输出" } }
},
"XB_VAEDecode": {
"display_name": "XB-BOX - VAE解码(原版优化)",
"inputs": {
"samples": { "name": "潜空间" },
"vae": { "name": "VAE" },
"cleanup": { "name": "执行前清理" }
},
"outputs": { "0": { "name": "图像" } }
},
"XB_VAEDecodeTiled": {
"display_name": "XB-BOX - VAE分块解码(原版优化)",
"inputs": {
"samples": { "name": "潜空间" },
"vae": { "name": "VAE" },
"tile_size": { "name": "分块大小" },
"overlap": { "name": "重叠" },
"temporal_size": { "name": "时间分块" },
"temporal_overlap": { "name": "时间重叠" },
"cleanup": { "name": "执行前清理" }
},
"outputs": { "0": { "name": "图像" } }
},
"XB_VAEDecodeTiledImage": {
"display_name": "XB-BOX - VAE解码(原版优化)",
"inputs": {
"samples": { "name": "潜空间" },
"vae": { "name": "VAE" },
"tile_size": { "name": "分块大小" },
"overlap": { "name": "重叠" },
"cleanup": { "name": "执行前清理" }
},
"outputs": { "0": { "name": "图像" } }
},
"XB_VAEEncode": {
"display_name": "XB-BOX - VAE编码(原版优化)",
"inputs": {
"pixels": { "name": "图像" },
"vae": { "name": "VAE" },
"cleanup": { "name": "执行前清理" }
},
"outputs": { "0": { "name": "潜空间" } }
},
"XB_VAEEncodeTiled": {
"display_name": "XB-BOX - VAE分块编码(原版优化)",
"inputs": {
"pixels": { "name": "图像" },
"vae": { "name": "VAE" },
"tile_size": { "name": "分块大小" },
"overlap": { "name": "重叠" },
"temporal_size": { "name": "时间分块" },
"temporal_overlap": { "name": "时间重叠" },
"cleanup": { "name": "执行前清理" }
},
"outputs": { "0": { "name": "潜空间" } }
},
"XB_VAEEncodeForInpaint": {
"display_name": "XB-BOX - VAE修补编码(原版优化)",
"inputs": {
"pixels": { "name": "图像" },
"vae": { "name": "VAE" },
"mask": { "name": "遮罩" },
"grow_mask_by": { "name": "遮罩扩展" },
"cleanup": { "name": "执行前清理" }
},
"outputs": { "0": { "name": "潜空间" } }
}
}
+273 -1530
View File
File diff suppressed because it is too large Load Diff
+34 -17
View File
@@ -122,35 +122,52 @@ class XB_SageAttentionAccelerator:
selected_cfg = configs[preset]
from sageattention import sageattn
_warned = {} # 去重:每个警告只打印一次
def attention_override_sage(func, q, k, v, heads, mask=None, attn_precision=None,
skip_reshape=False, skip_output_reshape=False, **kwargs):
in_dtype = v.dtype
if q.dtype == torch.float32 or k.dtype == torch.float32 or v.dtype == torch.float32:
q, k, v = q.to(torch.float16), k.to(torch.float16), v.to(torch.float16)
# 🛡️ 3D/4D 维度校验:非3D且非skip_reshape → 直接退回原生
if q.ndim != 3 and not skip_reshape:
return func(q, k, v, heads, mask=mask, attn_precision=attn_precision,
skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
# ── 在修改 q/k/v 物理形态之前,先完成所有断路检查 ──
if skip_reshape:
# q 已是 4D: (b, heads, n, dim_head),无需 reshape
b, _, _, dim_head = q.shape
tensor_layout = "HND"
else:
b, _, dim_head = q.shape
dim_head //= heads
# q 是 3D: (b, n, heads*dim_head),先提取维度信息,暂不 reshape
b, _, d = q.shape
dim_head = d // heads
tensor_layout = "NHD"
# 🛡️ 非2的幂Head Dim → 退回原生SDPA(此时 q 仍为原始形态)
if dim_head not in (16, 32, 64, 128, 256):
if "dim_head" not in _warned:
print("\033[93m[SageAttn]\033[0m: Unsupported head_dim ({}) — falling back to SDPA.".format(dim_head))
_warned["dim_head"] = True
return func(q, k, v, heads, mask=mask, attn_precision=attn_precision,
skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
# 🛡️ 有attn_mask → sage_config 路径不支持,退回原生SDPA(此时 q 仍为原始形态)
if mask is not None:
if "mask" not in _warned:
print("\033[93m[SageAttn]\033[0m: Attention mask detected — falling back to SDPA.")
_warned["mask"] = True
return func(q, k, v, heads, mask=mask, attn_precision=attn_precision,
skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
# ── 通过所有安全检查后,才进行 dtype 转换和 reshape ──
if q.dtype == torch.float32 or k.dtype == torch.float32 or v.dtype == torch.float32:
q, k, v = q.to(torch.float16), k.to(torch.float16), v.to(torch.float16)
if not skip_reshape:
q = q.view(b, -1, heads, dim_head)
k = k.view(b, -1, heads, dim_head)
v = v.view(b, -1, heads, dim_head)
tensor_layout = "NHD"
# 🛡️ 非2的幂Head Dim → 退回原生SDPA
if dim_head not in (16, 32, 64, 128, 256):
print("\033[93m[SageAttn]\033[0m: Unsupported head_dim ({}) — falling back to SDPA.".format(dim_head))
return func(q, k, v, heads, mask=mask, attn_precision=attn_precision,
skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
# 🛡️ 有attn_mask → sage_config 路径不支持,退回原生SDPA
if mask is not None:
print("\033[93m[SageAttn]\033[0m: Attention mask detected — falling back to SDPA.")
return func(q, k, v, heads, mask=mask, attn_precision=attn_precision,
skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
is_causal = kwargs.get("is_causal", False)
out = sageattn(q, k, v, is_causal=is_causal, tensor_layout=tensor_layout,
+523
View File
@@ -0,0 +1,523 @@
"""
XB_ToolBox 原版优化节点 (Vanilla Wrappers) — v1.0
=================================================
一字不差地调用 ComfyUI 官方节点,仅在其执行前插入分级显存清理。
彻底放弃自定义底层算法,完全依赖官方维护的采样/解码逻辑。
设计原则:
- 动态继承 INPUT_TYPES → 官方新增参数自动同步
- **kwargs 透传 → 参数名/数量永不过时
- 不做任何底层魔改 → 零维护成本
"""
import torch
import gc
import copy
import comfy.model_management as mm
import comfy.samplers
import nodes
# 安全导入: comfy_extras 在未来版本可能被移除
try:
from comfy_extras import nodes_custom_sampler as _ncs
except ImportError:
_ncs = None
# ═══════════════════════════════════════════════════════════════
# 统一的清理调度器
# ═══════════════════════════════════════════════════════════════
_CLEANUP_OPTIONS = [
"不做任何清理",
"单次缓存清理",
"卸载显存模型",
"卸载全量模型",
]
def _execute_cleanup(level: str, label: str = ""):
"""分级清理显存,在正式节点运行前强行制造显存真空区。
四级策略:
L0【不做任何清理】→ 透传
L1【单次缓存清理】→ soft_empty_cache + empty_cache,清除碎片
L2【卸载显存模型】→ unload_all_models,把模型退至内存
L3【卸载全量模型】→ L2 + cleanup_models + cleanup_models_gc + gc.collect + ipc_collect
"""
prefix = f"[XB_Wrapper{':'+label if label else ''}]"
if level == "不做任何清理":
return
print(f"{prefix} 🧹 执行 {level} ...", flush=True)
mm.soft_empty_cache()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if level in ("卸载显存模型", "卸载全量模型"):
mm.unload_all_models()
mm.soft_empty_cache()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if level == "卸载全量模型":
mm.cleanup_models()
mm.cleanup_models_gc()
gc.collect()
if torch.cuda.is_available() and hasattr(torch.cuda, "ipc_collect"):
torch.cuda.ipc_collect()
print(f"{prefix} ✅ 清理完成,显存环境纯净", flush=True)
# ═══════════════════════════════════════════════════════════════
# 采样器 Wrapper 节点
# ═══════════════════════════════════════════════════════════════
class XB_KSampler:
"""原版 KSampler 套皮 — 仅注入清理,参数零改动"""
@classmethod
def INPUT_TYPES(cls):
inputs = copy.deepcopy(nodes.KSampler.INPUT_TYPES())
inputs["required"]["cleanup"] = (_CLEANUP_OPTIONS, {"default": "不做任何清理"})
return inputs
RETURN_TYPES = nodes.KSampler.RETURN_TYPES
FUNCTION = "sample"
CATEGORY = "XB_ToolBox/原版优化"
def sample(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "采样器")
return nodes.KSampler().sample(**kwargs)
class XB_KSamplerAdvanced:
"""原版 KSamplerAdvanced 套皮 — 仅注入清理,参数零改动"""
@classmethod
def INPUT_TYPES(cls):
inputs = copy.deepcopy(nodes.KSamplerAdvanced.INPUT_TYPES())
inputs["required"]["cleanup"] = (_CLEANUP_OPTIONS, {"default": "不做任何清理"})
return inputs
RETURN_TYPES = nodes.KSamplerAdvanced.RETURN_TYPES
FUNCTION = "sample"
CATEGORY = "XB_ToolBox/原版优化"
def sample(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "高级采样器")
return nodes.KSamplerAdvanced().sample(**kwargs)
class XB_SamplerCustom:
"""原版 SamplerCustom 套皮 — 仅注入清理,参数零改动"""
@classmethod
def INPUT_TYPES(cls):
inputs = copy.deepcopy(_ncs.SamplerCustom.INPUT_TYPES())
inputs["required"]["cleanup"] = (_CLEANUP_OPTIONS, {"default": "不做任何清理"})
return inputs
RETURN_TYPES = _ncs.SamplerCustom.RETURN_TYPES
FUNCTION = "sample"
CATEGORY = "XB_ToolBox/原版优化"
def sample(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "自定义采样器")
return _ncs.SamplerCustom.sample(**kwargs)
class XB_SamplerCustomAdvanced:
"""原版 SamplerCustomAdvanced 套皮 — 仅注入清理,参数零改动"""
@classmethod
def INPUT_TYPES(cls):
inputs = copy.deepcopy(_ncs.SamplerCustomAdvanced.INPUT_TYPES())
inputs["required"]["cleanup"] = (_CLEANUP_OPTIONS, {"default": "不做任何清理"})
return inputs
RETURN_TYPES = _ncs.SamplerCustomAdvanced.RETURN_TYPES
FUNCTION = "sample"
CATEGORY = "XB_ToolBox/原版优化"
def sample(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "自定义高级采样器")
return _ncs.SamplerCustomAdvanced.sample(**kwargs)
# ═══════════════════════════════════════════════════════════════
# VAE 解码器 Wrapper 节点
# ═══════════════════════════════════════════════════════════════
class XB_VAEDecode:
"""原版 VAEDecode 套皮 — 仅注入清理,参数零改动"""
@classmethod
def INPUT_TYPES(cls):
inputs = copy.deepcopy(nodes.VAEDecode.INPUT_TYPES())
inputs["required"]["cleanup"] = (_CLEANUP_OPTIONS, {"default": "不做任何清理"})
return inputs
RETURN_TYPES = nodes.VAEDecode.RETURN_TYPES
FUNCTION = "decode"
CATEGORY = "XB_ToolBox/原版优化"
def decode(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "VAE解码")
# 🛡️ AMD MIOpen 防护:确保 latent 是连续内存
samples = kwargs.get("samples")
if isinstance(samples, dict):
lat = samples.get("samples")
if hasattr(lat, 'is_contiguous') and not lat.is_contiguous():
kwargs["samples"] = {**samples, "samples": lat.contiguous()}
return nodes.VAEDecode().decode(**kwargs)
class XB_VAEDecodeTiled:
"""原版 VAEDecodeTiled 套皮 — 仅注入清理,参数零改动"""
@classmethod
def INPUT_TYPES(cls):
inputs = copy.deepcopy(nodes.VAEDecodeTiled.INPUT_TYPES())
inputs["required"]["cleanup"] = (_CLEANUP_OPTIONS, {"default": "不做任何清理"})
return inputs
RETURN_TYPES = nodes.VAEDecodeTiled.RETURN_TYPES
FUNCTION = "decode"
CATEGORY = "XB_ToolBox/原版优化"
def decode(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "VAE分块解码")
return nodes.VAEDecodeTiled().decode(**kwargs)
class XB_VAEDecodeTiledImage:
"""原版 VAEDecodeTiled 图片专用套皮 — 仅空间分块,无时间维度"""
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"samples": ("LATENT",),
"vae": ("VAE",),
"tile_size": ("INT", {"default": 512, "min": 256, "max": 4096, "step": 64,
"tooltip": "空间分块大小(像素)"}),
"overlap": ("INT", {"default": 64, "min": 0, "max": 256, "step": 16,
"tooltip": "空间块重叠(像素)"}),
"cleanup": (_CLEANUP_OPTIONS, {"default": "不做任何清理"}),
}}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "decode"
CATEGORY = "XB_ToolBox/原版优化"
def decode(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "VAE分块解码(图片)")
# 透传给官方 VAEDecodeTiled,时间参数用默认值
kwargs.setdefault("temporal_size", 64)
kwargs.setdefault("temporal_overlap", 8)
return nodes.VAEDecodeTiled().decode(**kwargs)
class XB_VAEEncode:
"""原版 VAEEncode 套皮 — 仅注入清理,参数零改动"""
@classmethod
def INPUT_TYPES(cls):
inputs = copy.deepcopy(nodes.VAEEncode.INPUT_TYPES())
inputs["required"]["cleanup"] = (_CLEANUP_OPTIONS, {"default": "不做任何清理"})
return inputs
RETURN_TYPES = nodes.VAEEncode.RETURN_TYPES
FUNCTION = "encode"
CATEGORY = "XB_ToolBox/原版优化"
def encode(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "VAE编码")
return nodes.VAEEncode().encode(**kwargs)
class XB_VAEEncodeTiled:
"""原版 VAEEncodeTiled 套皮 — 仅注入清理,参数零改动"""
@classmethod
def INPUT_TYPES(cls):
inputs = copy.deepcopy(nodes.VAEEncodeTiled.INPUT_TYPES())
inputs["required"]["cleanup"] = (_CLEANUP_OPTIONS, {"default": "不做任何清理"})
return inputs
RETURN_TYPES = nodes.VAEEncodeTiled.RETURN_TYPES
FUNCTION = "encode"
CATEGORY = "XB_ToolBox/原版优化"
def encode(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "VAE分块编码")
return nodes.VAEEncodeTiled().encode(**kwargs)
class XB_VAEEncodeForInpaint:
"""原版 VAEEncodeForInpaint 套皮 — 仅注入清理,参数零改动"""
@classmethod
def INPUT_TYPES(cls):
inputs = copy.deepcopy(nodes.VAEEncodeForInpaint.INPUT_TYPES())
inputs["required"]["cleanup"] = (_CLEANUP_OPTIONS, {"default": "不做任何清理"})
return inputs
RETURN_TYPES = nodes.VAEEncodeForInpaint.RETURN_TYPES
FUNCTION = "encode"
CATEGORY = "XB_ToolBox/原版优化"
def encode(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "VAE修补编码")
return nodes.VAEEncodeForInpaint().encode(**kwargs)
# ═══════════════════════════════════════════════════════════════
# 别名类 — 匹配旧 ROCm 节点参数名,保证旧工作流连线兼容
# 清理选项统一使用新的四级体系
# ═══════════════════════════════════════════════════════════════
class _AliasKSampler:
"""别名: XB_ROCmKSampler — latent→latent_image, sampler→sampler_name"""
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"model": ("MODEL",),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1}),
"sampler": (comfy.samplers.KSampler.SAMPLERS,),
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"latent": ("LATENT",),
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"cleanup": (_CLEANUP_OPTIONS, {"default": "不做任何清理"}),
}}
RETURN_TYPES = ("LATENT",)
FUNCTION = "sample"
CATEGORY = "XB_ToolBox/原版优化"
def sample(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "采样器")
if "latent" in kwargs:
kwargs["latent_image"] = kwargs.pop("latent")
if "sampler" in kwargs:
kwargs["sampler_name"] = kwargs.pop("sampler")
return nodes.KSampler().sample(**kwargs)
class _AliasKSamplerAdvanced:
"""别名: XB_ROCmKSamplerAdvanced — latent→latent_image, sampler→sampler_name"""
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"model": ("MODEL",),
"add_noise": (["enable", "disable"], {}),
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1}),
"sampler": (comfy.samplers.KSampler.SAMPLERS,),
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"latent": ("LATENT",),
"start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
"end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
"return_with_leftover_noise": (["disable", "enable"], {}),
"cleanup": (_CLEANUP_OPTIONS, {"default": "不做任何清理"}),
}}
RETURN_TYPES = ("LATENT",)
FUNCTION = "sample"
CATEGORY = "XB_ToolBox/原版优化"
def sample(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "高级采样器")
if "latent" in kwargs:
kwargs["latent_image"] = kwargs.pop("latent")
if "sampler" in kwargs:
kwargs["sampler_name"] = kwargs.pop("sampler")
return nodes.KSamplerAdvanced().sample(**kwargs)
class _AliasSamplerCustom:
"""别名: XB_ROCmSamplerCustom — 参数名一致,仅加清理"""
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"model": ("MODEL",),
"add_noise": ("BOOLEAN", {"default": True}),
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1}),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"sampler": ("SAMPLER",),
"sigmas": ("SIGMAS",),
"latent_image": ("LATENT",),
"cleanup": (_CLEANUP_OPTIONS, {"default": "不做任何清理"}),
}}
RETURN_TYPES = ("LATENT", "LATENT")
RETURN_NAMES = ("output", "denoised_output")
FUNCTION = "sample"
CATEGORY = "XB_ToolBox/原版优化"
def sample(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "自定义采样器")
return _ncs.SamplerCustom.sample(**kwargs)
class _AliasSamplerCustomAdvanced:
"""别名: XB_ROCmSamplerCustomAdvanced — 参数名一致,仅加清理"""
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"noise": ("NOISE",),
"guider": ("GUIDER",),
"sampler": ("SAMPLER",),
"sigmas": ("SIGMAS",),
"latent_image": ("LATENT",),
"cleanup": (_CLEANUP_OPTIONS, {"default": "不做任何清理"}),
}}
RETURN_TYPES = ("LATENT", "LATENT")
RETURN_NAMES = ("output", "denoised_output")
FUNCTION = "sample"
CATEGORY = "XB_ToolBox/原版优化"
def sample(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "自定义高级采样器")
return _ncs.SamplerCustomAdvanced.sample(**kwargs)
class _AliasVAEDecode:
"""别名: XB_ROCmVAEDecode — tile→tile_size, 走图片分块解码"""
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"samples": ("LATENT",),
"vae": ("VAE",),
"tile": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 64}),
"overlap": ("INT", {"default": 0, "min": 0, "max": 256, "step": 16}),
"cleanup": (_CLEANUP_OPTIONS, {"default": "不做任何清理"}),
}}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "decode"
CATEGORY = "XB_ToolBox/原版优化"
def decode(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "VAE解码")
if "tile" in kwargs:
kwargs["tile_size"] = kwargs.pop("tile") if kwargs["tile"] > 0 else 512
kwargs.setdefault("temporal_size", 64)
kwargs.setdefault("temporal_overlap", 8)
return nodes.VAEDecodeTiled().decode(**kwargs)
class _AliasVAEDecodeTemporal:
"""别名: XB_ROCmVAEDecodeTemporal — tile→tile_size, t_tile→temporal_size, t_overlap→temporal_overlap"""
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"samples": ("LATENT",),
"vae": ("VAE",),
"tile": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 64}),
"overlap": ("INT", {"default": 0, "min": 0, "max": 256, "step": 16}),
"t_tile": ("INT", {"default": 0, "min": 0, "max": 1024, "step": 16}),
"t_overlap": ("INT", {"default": 0, "min": 0, "max": 128, "step": 8}),
"cleanup": (_CLEANUP_OPTIONS, {"default": "不做任何清理"}),
}}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "decode"
CATEGORY = "XB_ToolBox/原版优化"
def decode(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "VAE分块解码")
if "tile" in kwargs:
kwargs["tile_size"] = kwargs.pop("tile") if kwargs["tile"] > 0 else 512
if "t_tile" in kwargs:
kwargs["temporal_size"] = kwargs.pop("t_tile") if kwargs["t_tile"] > 0 else 64
if "t_overlap" in kwargs:
kwargs["temporal_overlap"] = kwargs.pop("t_overlap") if kwargs["t_overlap"] > 0 else 8
return nodes.VAEDecodeTiled().decode(**kwargs)
class _AliasVAEEncode:
"""别名: XB_ROCmVAEEncode — tile→tile_size"""
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"pixels": ("IMAGE",),
"vae": ("VAE",),
"tile": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 64}),
"overlap": ("INT", {"default": 0, "min": 0, "max": 256, "step": 16}),
"temporal_size": ("INT", {"default": 0, "min": 0, "max": 1024, "step": 16}),
"temporal_overlap": ("INT", {"default": 0, "min": 0, "max": 64, "step": 4}),
"cleanup": (_CLEANUP_OPTIONS, {"default": "不做任何清理"}),
}}
RETURN_TYPES = ("LATENT",)
FUNCTION = "encode"
CATEGORY = "XB_ToolBox/原版优化"
def encode(self, **kwargs):
cleanup = kwargs.pop("cleanup", "不做任何清理")
_execute_cleanup(cleanup, "VAE编码")
if "tile" in kwargs:
kwargs["tile_size"] = kwargs.pop("tile") if kwargs["tile"] > 0 else 512
return nodes.VAEEncodeTiled().encode(**kwargs)
# LTX 别名直接复用 Temporal 别名类
_AliasLTXVAEDecode = _AliasVAEDecodeTemporal
# ═══════════════════════════════════════════════════════════════
# 节点注册
# ═══════════════════════════════════════════════════════════════
NODE_CLASS_MAPPINGS = {
"XB_KSampler": XB_KSampler,
"XB_KSamplerAdvanced": XB_KSamplerAdvanced,
"XB_VAEDecode": XB_VAEDecode,
"XB_VAEDecodeTiled": XB_VAEDecodeTiled,
"XB_VAEDecodeTiledImage": XB_VAEDecodeTiledImage,
"XB_VAEEncode": XB_VAEEncode,
"XB_VAEEncodeTiled": XB_VAEEncodeTiled,
"XB_VAEEncodeForInpaint": XB_VAEEncodeForInpaint,
}
# 仅在 comfy_extras 可用时注册自定义采样器
if _ncs is not None:
NODE_CLASS_MAPPINGS["XB_SamplerCustom"] = XB_SamplerCustom
NODE_CLASS_MAPPINGS["XB_SamplerCustomAdvanced"] = XB_SamplerCustomAdvanced
NODE_DISPLAY_NAME_MAPPINGS = {
"XB_KSampler": "XB-采样器(原版优化)",
"XB_KSamplerAdvanced": "XB-高级采样器(原版优化)",
"XB_VAEDecode": "XB-VAE解码(原版优化)",
"XB_VAEDecodeTiled": "XB-VAE分块解码(原版优化)",
"XB_VAEDecodeTiledImage": "XB-VAE解码(原版优化)",
"XB_VAEEncode": "XB-VAE编码(原版优化)",
"XB_VAEEncodeTiled": "XB-VAE分块编码(原版优化)",
"XB_VAEEncodeForInpaint": "XB-VAE修补编码(原版优化)",
}
if _ncs is not None:
NODE_DISPLAY_NAME_MAPPINGS["XB_SamplerCustom"] = "XB-自定义采样器(原版优化)"
NODE_DISPLAY_NAME_MAPPINGS["XB_SamplerCustomAdvanced"] = "XB-自定义高级采样器(原版优化)"
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "XB_ToolBox"
description = "XB_ToolBox: An easy-to-use ComfyUI custom node suite that streamlines workflow generation. Featuring exclusive data-flow logic and pioneering experimental kernel optimizations dedicated to the native AMD GPU ecosystem."
version = "0.6.30"
version = "0.7.01"
license = { file = "LICENSE" }
dependencies = ["opencv-python", "easyocr"]