FearL0rd
6c9af0ff02
Updated the INPUT_TYPES definition to set "Vae Patched" as the default selection. Why: The "Original" VAE mode often causes black images or NaN errors in mixed-precision workflows (like Flux, SD3, or FP16 contexts) on certain hardware. Fix: By defaulting to "Vae Patched," the node now automatically handles precision casting (forcing Float32 when necessary) to ensure successful image decoding out of the box, while still allowing advanced users to switch back to "Original" if needed.
317 lines
14 KiB
Python
317 lines
14 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import types
|
|
import comfy.model_management
|
|
import gc
|
|
import sys
|
|
|
|
# --- GLOBAL ATTENTION CONTROLLER ---
|
|
try:
|
|
import xformers.ops
|
|
if not hasattr(xformers.ops, "original_memory_efficient_attention"):
|
|
xformers.ops.original_memory_efficient_attention = xformers.ops.memory_efficient_attention
|
|
|
|
def traffic_controlled_attention(query, key, value, attn_bias=None, p=0.0, scale=None):
|
|
device_type = query.device.type
|
|
original_dtype = query.dtype # Save for later restoration
|
|
|
|
# 1. HARDWARE CHECK (SM < 80)
|
|
use_fast_kernels = True
|
|
if device_type == "cuda":
|
|
try:
|
|
props = torch.cuda.get_device_properties(query.device)
|
|
sm_version = props.major * 10 + props.minor
|
|
if sm_version < 80:
|
|
use_fast_kernels = False
|
|
except:
|
|
pass
|
|
|
|
# 2. EXECUTION PATHS
|
|
if use_fast_kernels and device_type != "cpu":
|
|
return xformers.ops.original_memory_efficient_attention(query, key, value, attn_bias, p, scale)
|
|
|
|
# PATH B: Legacy / CPU Fallback
|
|
if query.dtype != torch.float32: query = query.to(torch.float32)
|
|
if key.dtype != torch.float32: key = key.to(torch.float32)
|
|
if value.dtype != torch.float32: value = value.to(torch.float32)
|
|
|
|
q = query.transpose(1, 2)
|
|
k = key.transpose(1, 2)
|
|
v = value.transpose(1, 2)
|
|
|
|
is_causal = False
|
|
attn_mask = None
|
|
if attn_bias is not None:
|
|
if isinstance(attn_bias, xformers.ops.LowerTriangularMask):
|
|
is_causal = True
|
|
attn_mask = None
|
|
elif isinstance(attn_bias, torch.Tensor):
|
|
attn_mask = attn_bias
|
|
if attn_mask.dtype != torch.float32 and attn_mask.dtype != torch.bool:
|
|
attn_mask = attn_mask.to(torch.float32)
|
|
|
|
out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=p, is_causal=is_causal, scale=scale)
|
|
|
|
# Restore Dtype (Fixes float != Half error)
|
|
if out.dtype != original_dtype:
|
|
out = out.to(original_dtype)
|
|
|
|
return out.transpose(1, 2)
|
|
|
|
xformers.ops.memory_efficient_attention = traffic_controlled_attention
|
|
except ImportError:
|
|
pass
|
|
# -----------------------------
|
|
|
|
def get_available_devices():
|
|
devices = []
|
|
if torch.cuda.is_available():
|
|
count = torch.cuda.device_count()
|
|
for i in range(count):
|
|
devices.append(f"cuda:{i}")
|
|
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
|
devices.append("mps")
|
|
if hasattr(torch, "xpu") and torch.xpu.is_available():
|
|
devices.append("xpu")
|
|
devices.append("cpu")
|
|
return devices
|
|
|
|
class AnyDeviceOffload:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
device_list = get_available_devices()
|
|
return {
|
|
"required": {
|
|
"target_device": (device_list, ),
|
|
# Swapped order: Vae Patched is now the first item (Default)
|
|
"vae_mode": (["Vae Patched", "Original"], ),
|
|
"keep_in_memory": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"model": ("MODEL",),
|
|
"clip": ("CLIP",),
|
|
"vae": ("VAE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
|
RETURN_NAMES = ("MODEL", "CLIP", "VAE")
|
|
FUNCTION = "offload_process"
|
|
CATEGORY = "utils/hardware"
|
|
|
|
def offload_process(self, target_device, vae_mode, keep_in_memory, model=None, clip=None, vae=None):
|
|
device_str = target_device.split(" ")[0]
|
|
try:
|
|
device = torch.device(device_str)
|
|
except:
|
|
device = torch.device("cpu")
|
|
|
|
is_cpu = device.type == 'cpu'
|
|
offload_target = device if keep_in_memory else torch.device("cpu")
|
|
|
|
print(f"[Offload] Target: {device} | Keep: {keep_in_memory}")
|
|
|
|
# --- 1. Handle MODEL (UNet/Transformer) ---
|
|
if model is not None:
|
|
target_model = None
|
|
if hasattr(model, "model"):
|
|
target_model = model.model
|
|
|
|
if target_model:
|
|
# A. Weight Conversion
|
|
if is_cpu:
|
|
target_model.to(dtype=torch.float32)
|
|
|
|
# B. FORCE DEVICE SYNC (Initial Move)
|
|
if keep_in_memory:
|
|
try:
|
|
target_model.to(device)
|
|
except:
|
|
pass
|
|
|
|
if hasattr(model, "patcher"):
|
|
model.patcher.load_device = device
|
|
model.patcher.offload_device = offload_target
|
|
model.patcher.current_device = device
|
|
|
|
# C. [NUCLEAR FIX] INSTALL RUNTIME DEVICE GUARD
|
|
if not hasattr(target_model, "is_device_guarded"):
|
|
print(" -> [System] Installing Runtime Device Guard...")
|
|
|
|
if hasattr(target_model, "diffusion_model"):
|
|
victim_model = target_model.diffusion_model
|
|
victim_model.original_forward_guard = victim_model.forward
|
|
|
|
def guarded_forward(self, *args, **kwargs):
|
|
# 1. Identify Input Device (Check ARGS and KWARGS)
|
|
target_run_device = None
|
|
|
|
# Scan positional args
|
|
for arg in args:
|
|
if isinstance(arg, torch.Tensor):
|
|
target_run_device = arg.device
|
|
break
|
|
|
|
# Scan keyword args if not found yet
|
|
if target_run_device is None:
|
|
for v in kwargs.values():
|
|
if isinstance(v, torch.Tensor):
|
|
target_run_device = v.device
|
|
break
|
|
|
|
# 2. Self-Healing Check
|
|
if target_run_device is not None:
|
|
try:
|
|
current_param_device = next(self.parameters()).device
|
|
if current_param_device != target_run_device:
|
|
# print(f"[Guard] Mismatch! {current_param_device} vs {target_run_device}. Moving...")
|
|
self.to(target_run_device)
|
|
except Exception as e:
|
|
pass # Suppress generic guard errors
|
|
|
|
# 3. CPU Float32 Casting
|
|
# [FIX] Ensure target_run_device exists before checking .type
|
|
if target_run_device is not None and target_run_device.type == "cpu":
|
|
new_args = []
|
|
for arg in args:
|
|
if isinstance(arg, torch.Tensor) and arg.dtype != torch.float32:
|
|
new_args.append(arg.to(torch.float32))
|
|
else:
|
|
new_args.append(arg)
|
|
new_kwargs = {}
|
|
for k, v in kwargs.items():
|
|
if isinstance(v, torch.Tensor) and v.dtype != torch.float32:
|
|
new_kwargs[k] = v.to(torch.float32)
|
|
else:
|
|
new_kwargs[k] = v
|
|
return self.original_forward_guard(*new_args, **new_kwargs)
|
|
else:
|
|
return self.original_forward_guard(*args, **kwargs)
|
|
|
|
victim_model.forward = types.MethodType(guarded_forward, victim_model)
|
|
target_model.is_device_guarded = True
|
|
|
|
# D. Update Internal Flags
|
|
model.load_device = device
|
|
model.offload_device = offload_target
|
|
if keep_in_memory:
|
|
model.current_device = device
|
|
|
|
# --- 2. Handle CLIP ---
|
|
if clip is not None:
|
|
if hasattr(clip, "cond_stage_model"):
|
|
target_clip = clip.cond_stage_model
|
|
if is_cpu:
|
|
target_clip.to(dtype=torch.float32)
|
|
|
|
if keep_in_memory:
|
|
try:
|
|
target_clip.to(device)
|
|
except:
|
|
pass
|
|
|
|
if hasattr(clip, "patcher"):
|
|
clip.patcher.load_device = device
|
|
clip.patcher.offload_device = offload_target
|
|
if keep_in_memory:
|
|
clip.patcher.current_device = device
|
|
|
|
# --- 3. Handle VAE ---
|
|
if vae is not None:
|
|
target_vae = None
|
|
if hasattr(vae, "first_stage_model"):
|
|
target_vae = vae.first_stage_model
|
|
elif hasattr(vae, "model"):
|
|
target_vae = vae.model
|
|
|
|
if target_vae:
|
|
if not hasattr(target_vae, "offload_node_state"):
|
|
target_vae.offload_node_state = {}
|
|
target_vae.offload_node_state['keep'] = keep_in_memory
|
|
target_vae.offload_node_state['device'] = device
|
|
target_vae.offload_node_state['is_cpu'] = is_cpu
|
|
target_vae.offload_node_state['model_ref'] = model
|
|
target_vae.offload_node_state['clip_ref'] = clip
|
|
|
|
if is_cpu or vae_mode == "Vae Patched":
|
|
target_vae.to(dtype=torch.float32)
|
|
|
|
if keep_in_memory:
|
|
try:
|
|
target_vae.to(device)
|
|
target_vae.eval()
|
|
except:
|
|
pass
|
|
|
|
# VAE Wrapper (Preserved)
|
|
if not hasattr(target_vae, "is_offload_patched"):
|
|
target_vae.original_decode_patched = target_vae.decode
|
|
|
|
def dynamic_kill_switch_wrapper(self, z, *args, **kwargs):
|
|
state = getattr(self, "offload_node_state", {})
|
|
should_keep = state.get('keep', True)
|
|
running_on_cpu = state.get('is_cpu', False)
|
|
|
|
try:
|
|
# 1. Type Safety
|
|
if running_on_cpu:
|
|
if z.dtype != torch.float32: z = z.to(torch.float32)
|
|
else:
|
|
dtype_target = next(self.parameters()).dtype
|
|
if z.dtype != dtype_target: z = z.to(dtype_target)
|
|
|
|
# 2. Device Safety
|
|
current_weight_device = next(self.parameters()).device
|
|
if z.device != current_weight_device:
|
|
z = z.to(current_weight_device)
|
|
|
|
return self.original_decode_patched(z, *args, **kwargs)
|
|
|
|
finally:
|
|
# 3. KILL SWITCH
|
|
if not should_keep:
|
|
self.to("cpu")
|
|
linked_model = state.get('model_ref')
|
|
linked_clip = state.get('clip_ref')
|
|
|
|
if linked_model and hasattr(linked_model, "model"):
|
|
linked_model.model.to("cpu")
|
|
linked_model.current_device = torch.device("cpu")
|
|
if hasattr(linked_model, "patcher"):
|
|
linked_model.patcher.current_device = torch.device("cpu")
|
|
|
|
if linked_clip and hasattr(linked_clip, "cond_stage_model"):
|
|
linked_clip.cond_stage_model.to("cpu")
|
|
if hasattr(linked_clip, "patcher"):
|
|
linked_clip.patcher.current_device = torch.device("cpu")
|
|
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
comfy.model_management.cleanup_models()
|
|
comfy.model_management.soft_empty_cache()
|
|
|
|
target_vae.decode = types.MethodType(dynamic_kill_switch_wrapper, target_vae)
|
|
target_vae.is_offload_patched = True
|
|
|
|
if hasattr(vae, "device") and keep_in_memory:
|
|
vae.device = device
|
|
|
|
if hasattr(vae, "patcher"):
|
|
vae.patcher.load_device = device
|
|
vae.patcher.offload_device = offload_target
|
|
if keep_in_memory:
|
|
vae.patcher.current_device = device
|
|
|
|
return (model, clip, vae)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AnyDeviceOffload": AnyDeviceOffload
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AnyDeviceOffload": "Offload Anything (GPU/CPU)"
|
|
}
|