nodes/capitan01R-Capitan-ConditioningEnhancer
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update enhancer.py enhanced memory management and changed name to "Capitan Enhancer ( Basic)"

Browse Source
This commit is contained in:
capitan01R
2026-01-15 03:19:13 -05:00
committed by GitHub
parent 268b403e89
commit b520a69dde
1 changed files with 48 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enhancer.py
+48 -21
View File
@@ -1,9 +1,15 @@
import torch
import comfy.model_management as mm
import gc
class ConditioningEnhancer:
@classmethod
def INPUT_TYPES(cls):
devices = ["auto", "cpu"]
if torch.cuda.is_available():
for i in range(torch.cuda.device_count()):
devices.append(f"cuda:{i}")
return {
"required": {
"conditioning": ("CONDITIONING", ),
@@ -11,41 +17,50 @@ class ConditioningEnhancer:
"normalize": ("BOOLEAN", {"default": True}),
"add_self_attention": ("BOOLEAN", {"default": False}),
"mlp_hidden_mult": ("INT", {"default": 2, "min": 1, "max": 100, "step": 1}),
"seed": ("INT", {"default": 42, "min": 0, "max": 2147483647, "step": 1}), # NEW: seed control
"seed": ("INT", {"default": 42, "min": 0, "max": 2147483647, "step": 1}),
"low_vram": ("BOOLEAN", {"default": False}),
"device": (devices, {"default": "auto"}),
}
}
RETURN_TYPES = ("CONDITIONING", )
FUNCTION = "enhance"
CATEGORY = "conditioning/enhance"
CATEGORY = "conditioning/enhance/capitan"
def enhance(self, conditioning, enhance_strength, normalize, add_self_attention, mlp_hidden_mult, seed):
def enhance(self, conditioning, enhance_strength, normalize, add_self_attention, mlp_hidden_mult, seed, low_vram, device):
if not conditioning:
return (conditioning, )
device = mm.get_torch_device()
if device == "auto":
device = mm.get_torch_device()
else:
device = torch.device(device)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
compute_dtype = torch.float16 if low_vram and device.type == "cuda" else torch.float32
enhanced = []
# NEW: Set seed for reproducible init
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
for emb_tensor, meta in conditioning:
emb = emb_tensor.to(device) # [B, seq, 2560]
orig_dtype = emb.dtype
for cond, meta in conditioning:
emb = cond.to(device, dtype=compute_dtype)
original_dtype = cond.dtype
if normalize:
emb = (emb - emb.mean(dim=-1, keepdim=True)) / (emb.std(dim=-1, keepdim=True) + 1e-6)
mean = emb.mean(dim=-1, keepdim=True)
std = emb.std(dim=-1, keepdim=True) + 1e-6
emb = (emb - mean) / std
dim = emb.shape[-1]
hidden_dim = dim * mlp_hidden_mult
mlp = torch.nn.Sequential(
torch.nn.Linear(dim, hidden_dim, device=device, dtype=emb.dtype),
torch.nn.Linear(dim, hidden_dim, device=device, dtype=compute_dtype),
torch.nn.GELU(),
torch.nn.Linear(hidden_dim, dim, device=device, dtype=emb.dtype)
).to(device)
torch.nn.Linear(hidden_dim, dim, device=device, dtype=compute_dtype)
)
torch.nn.init.kaiming_uniform_(mlp[0].weight, nonlinearity='relu')
torch.nn.init.zeros_(mlp[0].bias)
@@ -53,17 +68,29 @@ class ConditioningEnhancer:
torch.nn.init.zeros_(mlp[2].bias)
refined = mlp(emb)
blended = emb + enhance_strength * (refined - emb)
if add_self_attention:
attn = torch.nn.MultiheadAttention(embed_dim=dim, num_heads=8, batch_first=True, device=device, dtype=emb.dtype)
if add_self_attention and not low_vram:
attn = torch.nn.MultiheadAttention(
embed_dim=dim,
num_heads=8,
batch_first=True,
device=device,
dtype=compute_dtype
)
attn_out, _ = attn(blended, blended, blended)
blended = blended + 0.3 * attn_out
del attn, attn_out
enhanced.append((blended.to("cpu", dtype=orig_dtype), meta))
enhanced.append((blended.to("cpu", dtype=original_dtype), meta))
del mlp, refined, blended, emb
if device.type == "cuda":
torch.cuda.empty_cache()
gc.collect()
return (enhanced, )
NODE_CLASS_MAPPINGS = {"ConditioningEnhancer": ConditioningEnhancer}
NODE_DISPLAY_NAME_MAPPINGS = {"ConditioningEnhancer": "Conditioning Enhancer (Qwen/Z-Image)"}
NODE_DISPLAY_NAME_MAPPINGS = {"ConditioningEnhancer": "Capitan Conditioning Enhancer (basic)"}