FearL0rd
7283facab3
This update introduces a robust "Runtime Device Guard" and intelligent hardware detection to solve crashes on older GPUs and multi-GPU workflows. Example; # SM 8.0 = Ampere (RTX 30xx, A100) # SM 7.5 = Turing (RTX 20xx, T4) # SM 6.1 = Pascal (GTX 10xx)
316 lines
14 KiB
Python
316 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 (Preserves V17 Fixes) ---
|
|
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, ),
|
|
"vae_mode": (["Original", "Vae Patched"], ),
|
|
"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
|
|
|
|
# [FIX] 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)"
|
|
}
|