# 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()