364 lines
14 KiB
Python
364 lines
14 KiB
Python
# Comfyui扩展
|
|
import os
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import comfy.ops
|
|
from comfy.quant_ops import QuantizedTensor
|
|
import cmpext3
|
|
|
|
try:
|
|
from server import PromptServer
|
|
except Exception:
|
|
PromptServer = None
|
|
|
|
# =================================================================
|
|
# 1. 原始算子备份 (建立基准)
|
|
# =================================================================
|
|
# 我们在模块加载时就备份好原始的 PyTorch 函数
|
|
_ORIGINAL_OPS = {
|
|
# 基础运算
|
|
"linear": F.linear,
|
|
"conv2d": F.conv2d,
|
|
"conv_transpose2d": F.conv_transpose2d,
|
|
"bmm": torch.bmm,
|
|
|
|
# 归一化
|
|
"group_norm": F.group_norm,
|
|
"layer_norm": F.layer_norm,
|
|
|
|
# 激活函数
|
|
"silu": F.silu,
|
|
"gelu": F.gelu,
|
|
"softmax": F.softmax,
|
|
"mish": F.mish,
|
|
"softplus": F.softplus,
|
|
"softsign": F.softsign,
|
|
}
|
|
|
|
# =================================================================
|
|
# 2. 高健壮性 Wrappers (防止内存不连续导致的 Crash)
|
|
# =================================================================
|
|
|
|
def ensure_contiguous(tensor):
|
|
if tensor is None: return None
|
|
if not tensor.is_contiguous():
|
|
return tensor.contiguous()
|
|
return tensor
|
|
|
|
_CMP_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
|
|
_DEFAULT_CONFIG = {
|
|
"linear": False,
|
|
"conv2d": False,
|
|
"conv_transpose2d": False,
|
|
"bmm": False,
|
|
"group_norm": False,
|
|
"layer_norm": False,
|
|
"silu": False,
|
|
"gelu": False,
|
|
"softmax": False,
|
|
"mish": False,
|
|
"softplus": False,
|
|
"softsign": False,
|
|
"bf16_safe_accum": False,
|
|
"bf16_accum_mode": "auto",
|
|
}
|
|
_CURRENT_CONFIG = _DEFAULT_CONFIG.copy()
|
|
_PROMPT_RESET_REGISTERED = False
|
|
_BF16_SAFE_MODEL_HINTS = ("lumina", "flux", "hidream")
|
|
|
|
def is_plain_tensor(tensor):
|
|
return isinstance(tensor, torch.Tensor) and not isinstance(tensor, QuantizedTensor)
|
|
|
|
def can_use_cmpext3_linear(input, weight, bias):
|
|
if not is_plain_tensor(input) or not is_plain_tensor(weight):
|
|
return False
|
|
if input.dim() < 2 or weight.dim() != 2:
|
|
return False
|
|
if not input.is_cuda or not weight.is_cuda:
|
|
return False
|
|
if input.device != weight.device:
|
|
return False
|
|
if input.dtype not in _CMP_DTYPES or weight.dtype != input.dtype:
|
|
return False
|
|
if input.shape[-1] != weight.shape[1]:
|
|
return False
|
|
if bias is not None:
|
|
if not is_plain_tensor(bias):
|
|
return False
|
|
if not bias.is_cuda or bias.device != input.device or bias.dtype != input.dtype:
|
|
return False
|
|
if bias.dim() != 1 or bias.shape[0] != weight.shape[0]:
|
|
return False
|
|
return True
|
|
|
|
def op_enabled(name):
|
|
return bool(_CURRENT_CONFIG.get(name, False))
|
|
|
|
def describe_model(model):
|
|
parts = []
|
|
seen = set()
|
|
queue = [model]
|
|
for _ in range(16):
|
|
if not queue:
|
|
break
|
|
obj = queue.pop(0)
|
|
if obj is None or id(obj) in seen:
|
|
continue
|
|
seen.add(id(obj))
|
|
parts.append(obj.__class__.__name__)
|
|
for attr in ("model", "diffusion_model", "model_config", "inner_model", "model_type", "unet_config"):
|
|
value = getattr(obj, attr, None)
|
|
if isinstance(value, dict):
|
|
parts.extend(str(item) for pair in value.items() for item in pair[:2])
|
|
continue
|
|
if isinstance(value, (str, int, float, bool)):
|
|
parts.append(str(value))
|
|
continue
|
|
queue.append(getattr(obj, attr, None))
|
|
return " ".join(parts).lower()
|
|
|
|
def infer_bf16_safe_accum(model, mode="auto"):
|
|
env_mode = os.environ.get("CMPEXT3_BF16_SAFE_MODE")
|
|
if env_mode is not None:
|
|
mode = env_mode
|
|
mode = str(mode).strip().lower()
|
|
if mode in ("1", "true", "yes", "on", "safe"):
|
|
return True
|
|
if mode in ("0", "false", "no", "off", "fast"):
|
|
return False
|
|
|
|
desc = describe_model(model)
|
|
return any(hint in desc for hint in _BF16_SAFE_MODEL_HINTS)
|
|
|
|
def use_bf16_safe_accum(input, weight):
|
|
if input.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16:
|
|
return False
|
|
return bool(_CURRENT_CONFIG.get("bf16_safe_accum", False))
|
|
|
|
# --- Linear & Conv ---
|
|
def cmpext3_linear_wrapper(input, weight, bias=None):
|
|
if not op_enabled("linear"):
|
|
return _ORIGINAL_OPS["linear"](input, weight, bias)
|
|
if not can_use_cmpext3_linear(input, weight, bias):
|
|
return _ORIGINAL_OPS["linear"](input, weight, bias)
|
|
|
|
is_multidim = input.dim() > 2
|
|
original_shape = input.shape
|
|
x = ensure_contiguous(input)
|
|
w = ensure_contiguous(weight)
|
|
if is_multidim: x = x.reshape(-1, original_shape[-1])
|
|
|
|
b = ensure_contiguous(bias)
|
|
|
|
output = cmpext3.linear(x, w, b, use_bf16_safe_accum(x, w))
|
|
if is_multidim: output = output.view(original_shape[:-1] + (output.shape[-1],))
|
|
return output
|
|
|
|
def cmpext3_conv2d_wrapper(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
|
|
if not op_enabled("conv2d"):
|
|
return _ORIGINAL_OPS["conv2d"](input, weight, bias, stride, padding, dilation, groups)
|
|
b = ensure_contiguous(bias) if bias is not None else torch.zeros(weight.shape[0], device=input.device, dtype=input.dtype)
|
|
return cmpext3.conv2d(ensure_contiguous(input), ensure_contiguous(weight), b, stride, padding, dilation, groups)
|
|
|
|
def cmpext3_conv_transpose2d_wrapper(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1):
|
|
if not op_enabled("conv_transpose2d"):
|
|
return _ORIGINAL_OPS["conv_transpose2d"](input, weight, bias, stride, padding, output_padding, groups, dilation)
|
|
b = ensure_contiguous(bias) # 允许为 None，视 C++ 实现而定
|
|
return cmpext3.conv_transpose2d(ensure_contiguous(input), ensure_contiguous(weight), b, stride, padding, output_padding, dilation, groups)
|
|
|
|
def cmpext3_bmm_wrapper(input, mat2, *, out=None):
|
|
if not op_enabled("bmm"):
|
|
return _ORIGINAL_OPS["bmm"](input, mat2, out=out)
|
|
return cmpext3.bmm(ensure_contiguous(input), ensure_contiguous(mat2))
|
|
|
|
# --- Norms ---
|
|
def cmpext3_group_norm_wrapper(input, num_groups, weight=None, bias=None, eps=1e-5):
|
|
if not op_enabled("group_norm"):
|
|
return _ORIGINAL_OPS["group_norm"](input, num_groups, weight, bias, eps)
|
|
return cmpext3.group_norm(ensure_contiguous(input), num_groups, ensure_contiguous(weight), ensure_contiguous(bias), eps)
|
|
|
|
def cmpext3_layer_norm_wrapper(input, normalized_shape, weight=None, bias=None, eps=1e-5):
|
|
if not op_enabled("layer_norm"):
|
|
return _ORIGINAL_OPS["layer_norm"](input, normalized_shape, weight, bias, eps)
|
|
return cmpext3.layer_norm(ensure_contiguous(input), normalized_shape, ensure_contiguous(weight), ensure_contiguous(bias), eps)
|
|
|
|
# --- Activations (重点排查区域) ---
|
|
def cmpext3_silu_wrapper(input, inplace=False):
|
|
if not op_enabled("silu"):
|
|
return _ORIGINAL_OPS["silu"](input, inplace=inplace)
|
|
return cmpext3.silu(ensure_contiguous(input))
|
|
|
|
def cmpext3_gelu_wrapper(input, approximate='none'):
|
|
if not op_enabled("gelu"):
|
|
return _ORIGINAL_OPS["gelu"](input, approximate=approximate)
|
|
return cmpext3.gelu(ensure_contiguous(input))
|
|
|
|
def cmpext3_softmax_wrapper(input, dim=None, _stacklevel=3, dtype=None):
|
|
if not op_enabled("softmax"):
|
|
return _ORIGINAL_OPS["softmax"](input, dim=dim, _stacklevel=_stacklevel, dtype=dtype)
|
|
if dtype is not None: input = input.to(dtype=dtype)
|
|
if dim is None: dim = -1
|
|
return cmpext3.softmax(ensure_contiguous(input), dim, False)
|
|
|
|
def cmpext3_mish_wrapper(input, inplace=False):
|
|
if not op_enabled("mish"):
|
|
return _ORIGINAL_OPS["mish"](input, inplace=inplace)
|
|
return cmpext3.mish(ensure_contiguous(input))
|
|
|
|
def cmpext3_softplus_wrapper(input, beta=1, threshold=20):
|
|
if not op_enabled("softplus"):
|
|
return _ORIGINAL_OPS["softplus"](input, beta=beta, threshold=threshold)
|
|
return cmpext3.softplus(ensure_contiguous(input), beta, threshold)
|
|
|
|
def cmpext3_softsign_wrapper(input):
|
|
if not op_enabled("softsign"):
|
|
return _ORIGINAL_OPS["softsign"](input)
|
|
return cmpext3.softsign(ensure_contiguous(input))
|
|
|
|
|
|
# =================================================================
|
|
# 3. 核心 Patch 管理器
|
|
# =================================================================
|
|
|
|
def apply_patch(target_obj, target_name, original_func, custom_wrapper, enable):
|
|
"""
|
|
通用 Patch 函数：
|
|
如果 enable 为 True，则将 target_obj.target_name 指向 custom_wrapper
|
|
如果 enable 为 False，则将 target_obj.target_name 还原为 original_func
|
|
"""
|
|
current_func = getattr(target_obj, target_name)
|
|
|
|
if enable:
|
|
if current_func != custom_wrapper:
|
|
setattr(target_obj, target_name, custom_wrapper)
|
|
# print(f" [ON] {target_name}")
|
|
else:
|
|
if current_func != original_func:
|
|
setattr(target_obj, target_name, original_func)
|
|
# print(f" [OFF] {target_name}")
|
|
|
|
def update_global_patches(config):
|
|
global _CURRENT_CONFIG
|
|
_CURRENT_CONFIG = _DEFAULT_CONFIG.copy()
|
|
_CURRENT_CONFIG.update(config)
|
|
|
|
print(f"\n[CmpExt3 Debugger] Syncing operators...")
|
|
|
|
# 1. 基础运算
|
|
apply_patch(F, "linear", _ORIGINAL_OPS["linear"], cmpext3_linear_wrapper, _CURRENT_CONFIG['linear'])
|
|
apply_patch(F, "conv2d", _ORIGINAL_OPS["conv2d"], cmpext3_conv2d_wrapper, _CURRENT_CONFIG['conv2d'])
|
|
apply_patch(F, "conv_transpose2d", _ORIGINAL_OPS["conv_transpose2d"], cmpext3_conv_transpose2d_wrapper, _CURRENT_CONFIG['conv_transpose2d'])
|
|
apply_patch(torch, "bmm", _ORIGINAL_OPS["bmm"], cmpext3_bmm_wrapper, _CURRENT_CONFIG['bmm'])
|
|
|
|
# 2. 归一化
|
|
apply_patch(F, "group_norm", _ORIGINAL_OPS["group_norm"], cmpext3_group_norm_wrapper, _CURRENT_CONFIG['group_norm'])
|
|
apply_patch(F, "layer_norm", _ORIGINAL_OPS["layer_norm"], cmpext3_layer_norm_wrapper, _CURRENT_CONFIG['layer_norm'])
|
|
|
|
# 3. 激活函数 (细粒度)
|
|
apply_patch(F, "silu", _ORIGINAL_OPS["silu"], cmpext3_silu_wrapper, _CURRENT_CONFIG['silu'])
|
|
apply_patch(F, "gelu", _ORIGINAL_OPS["gelu"], cmpext3_gelu_wrapper, _CURRENT_CONFIG['gelu'])
|
|
apply_patch(F, "softmax", _ORIGINAL_OPS["softmax"], cmpext3_softmax_wrapper, _CURRENT_CONFIG['softmax'])
|
|
apply_patch(F, "mish", _ORIGINAL_OPS["mish"], cmpext3_mish_wrapper, _CURRENT_CONFIG['mish'])
|
|
apply_patch(F, "softplus", _ORIGINAL_OPS["softplus"], cmpext3_softplus_wrapper, _CURRENT_CONFIG['softplus'])
|
|
apply_patch(F, "softsign", _ORIGINAL_OPS["softsign"], cmpext3_softsign_wrapper, _CURRENT_CONFIG['softsign'])
|
|
|
|
# 特殊处理：torch.softmax 别名
|
|
if hasattr(torch, 'softmax'):
|
|
apply_patch(torch, "softmax", _ORIGINAL_OPS["softmax"], cmpext3_softmax_wrapper, _CURRENT_CONFIG['softmax'])
|
|
|
|
print("[CmpExt3 Debugger] Update complete.")
|
|
|
|
def reset_global_patches_on_prompt(json_data=None):
|
|
update_global_patches(_DEFAULT_CONFIG)
|
|
return json_data
|
|
|
|
def register_prompt_reset():
|
|
global _PROMPT_RESET_REGISTERED
|
|
if _PROMPT_RESET_REGISTERED or PromptServer is None:
|
|
return
|
|
instance = getattr(PromptServer, "instance", None)
|
|
if instance is None:
|
|
return
|
|
instance.add_on_prompt_handler(reset_global_patches_on_prompt)
|
|
_PROMPT_RESET_REGISTERED = True
|
|
|
|
|
|
# =================================================================
|
|
# 4. ComfyUI 节点定义 (Control Panel)
|
|
# =================================================================
|
|
|
|
class CmpExt3ControlPanelNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
# --- Core ---
|
|
"enable_linear": ("BOOLEAN", {"default": True, "label": "Ops: Linear"}),
|
|
"enable_conv2d": ("BOOLEAN", {"default": False, "label": "Ops: Conv2d"}),
|
|
"enable_conv_transpose": ("BOOLEAN", {"default": False, "label": "Ops: ConvTranspose2d (VAE Upscale)"}),
|
|
"enable_bmm": ("BOOLEAN", {"default": True, "label": "Ops: BMM (Attention)"}),
|
|
|
|
# --- Norms ---
|
|
"enable_group_norm": ("BOOLEAN", {"default": True, "label": "Norm: GroupNorm (VAE)"}),
|
|
"enable_layer_norm": ("BOOLEAN", {"default": True, "label": "Norm: LayerNorm (CLIP/Transformer)"}),
|
|
|
|
# --- Activations (Suspects) ---
|
|
"enable_silu": ("BOOLEAN", {"default": True, "label": "Act: SiLU / Swish (VAE Main Suspect!)"}),
|
|
"enable_gelu": ("BOOLEAN", {"default": True, "label": "Act: GELU (CLIP/UNet)"}),
|
|
"enable_softmax": ("BOOLEAN", {"default": True, "label": "Act: Softmax (Attention)"}),
|
|
"enable_mish": ("BOOLEAN", {"default": False, "label": "Act: Mish"}),
|
|
"enable_softplus": ("BOOLEAN", {"default": False, "label": "Act: Softplus"}),
|
|
"enable_softsign": ("BOOLEAN", {"default": False, "label": "Act: Softsign"}),
|
|
"bf16_accum_mode": (["auto", "fast", "safe"], {"default": "auto", "label": "BF16 Accum"}),
|
|
},
|
|
"optional": {
|
|
"model": ("MODEL",),
|
|
"clip": ("CLIP",),
|
|
"vae": ("VAE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
|
RETURN_NAMES = ("model", "clip", "vae")
|
|
FUNCTION = "execute_patch"
|
|
CATEGORY = "eastmoe/Comfy-CMP"
|
|
|
|
def execute_patch(self,
|
|
enable_linear, enable_conv2d, enable_conv_transpose, enable_bmm,
|
|
enable_group_norm, enable_layer_norm,
|
|
enable_silu, enable_gelu, enable_softmax, enable_mish, enable_softplus, enable_softsign, bf16_accum_mode="auto",
|
|
model=None, clip=None, vae=None):
|
|
|
|
config = {
|
|
"linear": enable_linear,
|
|
"conv2d": enable_conv2d,
|
|
"conv_transpose2d": enable_conv_transpose,
|
|
"bmm": enable_bmm,
|
|
"group_norm": enable_group_norm,
|
|
"layer_norm": enable_layer_norm,
|
|
"silu": enable_silu,
|
|
"gelu": enable_gelu,
|
|
"softmax": enable_softmax,
|
|
"mish": enable_mish,
|
|
"softplus": enable_softplus,
|
|
"softsign": enable_softsign,
|
|
"bf16_accum_mode": bf16_accum_mode,
|
|
"bf16_safe_accum": infer_bf16_safe_accum(model, bf16_accum_mode),
|
|
}
|
|
|
|
register_prompt_reset()
|
|
update_global_patches(config)
|
|
|
|
return (model, clip, vae)
|
|
|
|
# 注册节点
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CmpExt3ControlPanel": CmpExt3ControlPanelNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"CmpExt3ControlPanel": "CmpExt3 Full Control Panel"
|
|
}
|
|
|
|
register_prompt_reset()
|