This repository has been archived on 2026-08-19. You can view files and clone it. You cannot open issues or pull requests or push a commit.
Files
2026-08-20 00:45:43 +02:00

632 lines
24 KiB
Python

"""
BC-250 gfx1010 Comprehensive Monkey-Patch v10
1. Replaces torch.softmax with manual implementation (VGPR overflow fix)
2. Replaces SDPA with manual implementation
3. Patches GGUF cast_bias_weight to dequant on CPU (avoids GPU page-fault hangs)
4. Pre-clones mmap'd tensor data before GPU transfer (XNACK workaround)
5. Pre-warms GPU context and caching allocator
6. Forces text encoder to CPU (memory constraint)
7. VAE decode on CPU float32 (bypasses GPU managed memory issues)
8. Sets torch threads to all CPU cores (faster CPU ops + VAE decode)
9. Caches VAE model on CPU (avoids reload each generation)
10. Startup preloading: submits warmup prompt to preload all models on boot
v10 changes:
- Background warmup thread submits 64x64 @ 1 step prompt after server starts
- All models (CLIP, UNET, VAE) preloaded before user interaction
- Models configurable via BC250_PRELOAD_* env vars
BC-250 APU / gfx1010: XNACK disabled, no GPU page fault handling.
GPU copy shader hangs on non-resident pages (mmap'd or swapped).
Place in ComfyUI root and import as first line of main.py.
"""
import torch
import torch.nn.functional as F
import os
import sys
import gc
import logging
import threading
import json
import time as _time
import threading
import json
import time as _time
logger = logging.getLogger(__name__)
# === THREAD CONFIGURATION ===
# BC-250 has 12 threads (6C/12T Zen2). Use all for CPU-heavy work (VAE, CLIP, dequant).
_NUM_THREADS = int(os.environ.get("BC250_NUM_THREADS", str(os.cpu_count() or 12)))
torch.set_num_threads(_NUM_THREADS)
# Note: set_num_interop_threads must be called before any parallel op, skip to avoid deadlock
logger.warning(f"[BC-250] Torch threads: intra-op={_NUM_THREADS}")
SAFE_SOFTMAX_THRESHOLD = int(os.environ.get("BC250_SOFTMAX_THRESHOLD", "4096"))
_original_softmax = torch.nn.functional.softmax
_original_tensor_softmax = torch.Tensor.softmax
_original_sdpa = torch.nn.functional.scaled_dot_product_attention
# === MMAP PRE-CLONE PATCH ===
_original_module_apply = torch.nn.Module._apply
def _bc250_safe_apply(self, fn, recurse=True):
"""Pre-clone mmap'd CPU tensor data before GPU transfer to avoid XNACK hangs."""
for key, param in self._parameters.items():
if param is not None and param.device.type == 'cpu':
param.data = param.data.clone()
for key, buf in self._buffers.items():
if buf is not None and buf.device.type == 'cpu':
self._buffers[key] = buf.clone()
return _original_module_apply(self, fn, recurse)
# === SOFTMAX PATCH ===
def _safe_softmax_impl(input, dim=-1):
x_max = input.max(dim=dim, keepdim=True).values
exp_x = torch.exp(input - x_max)
return exp_x / exp_x.sum(dim=dim, keepdim=True)
def patched_softmax(input, dim=None, _stacklevel=3, dtype=None):
if dim is None:
dim = -1
if dtype is not None:
input = input.to(dtype)
if input.is_cuda and input.shape[dim] > SAFE_SOFTMAX_THRESHOLD:
if not getattr(patched_softmax, '_logged', False):
logger.warning(f"[BC-250] Manual F.softmax triggered: shape={list(input.shape)}, dim={dim}, threshold={SAFE_SOFTMAX_THRESHOLD}")
patched_softmax._logged = True
return _safe_softmax_impl(input, dim)
return _original_softmax(input, dim=dim)
def patched_tensor_softmax(self, dim=-1, dtype=None):
if dtype is not None:
self = self.to(dtype)
if self.is_cuda and self.shape[dim] > SAFE_SOFTMAX_THRESHOLD:
if not getattr(patched_tensor_softmax, '_logged', False):
logger.warning(f"[BC-250] Manual softmax triggered: shape={list(self.shape)}, dim={dim}, threshold={SAFE_SOFTMAX_THRESHOLD}")
patched_tensor_softmax._logged = True
return _safe_softmax_impl(self, dim)
return _original_tensor_softmax(self, dim=dim)
def _safe_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
L, S = query.size(-2), key.size(-2)
if scale is None:
scale = query.size(-1) ** -0.5
attn_weight = torch.matmul(query, key.transpose(-2, -1)) * scale
if is_causal:
causal_mask = torch.triu(torch.ones(L, S, dtype=torch.bool, device=query.device), diagonal=1)
attn_weight = attn_weight.masked_fill(causal_mask, float('-inf'))
if attn_mask is not None:
if attn_mask.dtype == torch.bool:
attn_weight = attn_weight.masked_fill(~attn_mask, float('-inf'))
else:
attn_weight = attn_weight + attn_mask
attn_weight = _safe_softmax_impl(attn_weight, dim=-1)
if dropout_p > 0.0:
attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p)
return torch.matmul(attn_weight, value)
def patched_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
S = key.size(-2)
if query.is_cuda and S > SAFE_SOFTMAX_THRESHOLD:
if not getattr(patched_sdpa, '_logged', False):
logger.warning(f"[BC-250] Manual SDPA triggered: Q={list(query.shape)}, K={list(key.shape)}, S={S}, threshold={SAFE_SOFTMAX_THRESHOLD}")
patched_sdpa._logged = True
return _safe_sdpa(query, key, value, attn_mask=attn_mask,
dropout_p=dropout_p, is_causal=is_causal, scale=scale)
return _original_sdpa(query, key, value, attn_mask=attn_mask,
dropout_p=dropout_p, is_causal=is_causal, scale=scale)
# === GGUF CPU-DEQUANT PATCH (cast_bias_weight override) ===
_gguf_patched = False
def _try_patch_gguf():
"""Patch GGMLLayer.cast_bias_weight to dequant on CPU, send floats to GPU."""
global _gguf_patched
if _gguf_patched:
return True
ops_mod = None
dequant_mod = None
for name, mod in sys.modules.items():
if mod is None:
continue
if name.endswith('.ops') and 'GGUF' in name:
ops_mod = mod
if name.endswith('.dequant') and 'GGUF' in name:
dequant_mod = mod
if ops_mod is None or dequant_mod is None:
return False
GGMLLayer = getattr(ops_mod, 'GGMLLayer', None)
is_quantized_fn = getattr(dequant_mod, 'is_quantized', None)
if GGMLLayer is None or is_quantized_fn is None:
return False
def _bc250_cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
"""Dequant on CPU, send float results to GPU.
Cannot use .to(device) on quantized GGUF tensors from mmap'd files
(GPU copy shader hangs on non-resident pages, XNACK disabled).
Dequant to float on CPU, then transfer dequantized float to GPU.
"""
import comfy.model_management
import comfy.ops
if input is not None:
if dtype is None:
dtype = getattr(input, "dtype", torch.float32)
if bias_dtype is None:
bias_dtype = dtype
if device is None:
device = input.device
non_blocking = comfy.model_management.device_supports_non_blocking(device)
bias = None
if s.bias is not None:
if is_quantized_fn(s.bias):
bias = s.get_weight(s.bias, bias_dtype)
else:
bias = s.get_weight(s.bias.to(device), bias_dtype)
bias = comfy.ops.cast_to(bias, bias_dtype, device, non_blocking=non_blocking, copy=False)
if is_quantized_fn(s.weight):
weight = s.get_weight(s.weight, dtype)
else:
weight = s.get_weight(s.weight.to(device), dtype)
weight = comfy.ops.cast_to(weight, dtype, device, non_blocking=non_blocking, copy=False)
return weight, bias
GGMLLayer.cast_bias_weight = _bc250_cast_bias_weight
_gguf_patched = True
logger.warning("[BC-250] GGUF cast_bias_weight patched (CPU dequant)")
return True
# === IMPORT HOOK for deferred GGUF patching ===
class _GGUFImportWatcher:
def __init__(self):
self.done = False
def find_module(self, fullname, path=None):
if self.done:
return None
if 'GGUF' in fullname and ('dequant' in fullname or 'ops' in fullname):
return self
return None
def load_module(self, fullname):
if self in sys.meta_path:
sys.meta_path.remove(self)
try:
import importlib
mod = importlib.import_module(fullname)
finally:
if self not in sys.meta_path:
sys.meta_path.insert(0, self)
if _try_patch_gguf():
self.done = True
return mod
# === TEXT ENCODER CPU PATCH ===
_te_patched = False
def _try_patch_text_encoder_device():
global _te_patched
if _te_patched:
return True
mm = sys.modules.get('comfy.model_management')
if mm is None:
return False
mm.text_encoder_device = lambda: torch.device("cpu")
mm.text_encoder_offload_device = lambda: torch.device("cpu")
_te_patched = True
logger.warning("[BC-250] Text encoder forced to CPU (memory constraint)")
return True
# === VAE CPU FLOAT32 DECODE PATCH ===
# Decode VAE on CPU using float32 (not fp16). fp16 on CPU is emulated (10x slower).
# Cannot use GPU because UNet managed memory blocks new GPU allocations (XNACK disabled).
# 320MB VAE at float32 = 640MB RAM. For 256x256: ~2-3 min on 12-thread CPU.
_vae_patched = False
_vae_cached = False # Track whether VAE is already loaded to CPU float32
def _try_patch_vae_cpu():
"""Patch comfy.sd.VAE to decode on CPU with float32, with persistent caching."""
global _vae_patched
if _vae_patched:
return True
sd_mod = sys.modules.get('comfy.sd')
if sd_mod is None:
return False
VAE = getattr(sd_mod, 'VAE', None)
if VAE is None:
return False
_original_vae_encode = getattr(VAE, 'encode', None)
def _ensure_vae_on_cpu_f32(self):
"""Move VAE to CPU float32 once, then keep it cached."""
global _vae_cached
if not _vae_cached or next(self.first_stage_model.parameters()).dtype != torch.float32:
logger.warning("[BC-250] Loading VAE to CPU float32 (will stay cached)")
self.first_stage_model.to(torch.float32).to(torch.device("cpu"))
self.first_stage_model.eval()
_vae_cached = True
# Prevent ComfyUI model_management from offloading the VAE
self.disable_offload = True
def _bc250_vae_decode(self, samples_in, vae_options={}):
"""CPU float32 VAE decode — bypasses GPU managed memory entirely.
The UNet (5032MB managed memory) blocks new GPU allocations
when its pages are swapped by the OS (XNACK disabled on gfx1010).
float32 on CPU is ~5x faster than fp16 (which requires emulation).
VAE stays cached on CPU after first load — no re-conversion needed.
"""
import time
t0 = time.time()
logger.warning("[BC-250] VAE decode: CPU float32 (cached)")
self.throw_exception_if_invalid()
if self.latent_dim == 2 and samples_in.ndim == 5:
samples_in = samples_in[:, :, 0]
cpu = torch.device("cpu")
_ensure_vae_on_cpu_f32(self)
pixel_samples = None
with torch.no_grad():
for x in range(samples_in.shape[0]):
sample = samples_in[x:x+1].to(torch.float32)
decoded = self.first_stage_model.decode(sample, **vae_options)
# Squeeze temporal dim for 3D video autoencoders (single image)
if decoded.ndim == 5:
decoded = decoded[:, :, 0]
out = self.process_output(decoded.float())
if pixel_samples is None:
pixel_samples = torch.empty(
(samples_in.shape[0],) + tuple(out.shape[1:]),
device=cpu
)
pixel_samples[x:x+1] = out
del decoded, sample
# NCHW → NHWC (same as original ComfyUI VAE.decode line 977)
pixel_samples = pixel_samples.movedim(1, -1)
elapsed = time.time() - t0
logger.warning(f"[BC-250] VAE decode complete in {elapsed:.1f}s")
return pixel_samples
VAE.decode = _bc250_vae_decode
if _original_vae_encode is not None:
def _bc250_vae_encode(self, pixel_samples):
"""CPU float32 VAE encode (cached)."""
import time
t0 = time.time()
logger.warning("[BC-250] VAE encode: CPU float32 (cached)")
self.throw_exception_if_invalid()
_ensure_vae_on_cpu_f32(self)
with torch.no_grad():
pixels_in = self.process_input(pixel_samples).to(torch.float32)
result = self.first_stage_model.encode(pixels_in).float()
elapsed = time.time() - t0
logger.warning(f"[BC-250] VAE encode complete in {elapsed:.1f}s")
return result
VAE.encode = _bc250_vae_encode
_vae_patched = True
logger.warning("[BC-250] VAE patched: CPU float32 decode/encode with caching (bypass GPU managed memory)")
return True
class _SDModuleWatcher:
"""Patches comfy.sd.VAE after it's imported."""
def __init__(self):
self.done = False
def find_module(self, fullname, path=None):
if self.done:
return None
if fullname == 'comfy.sd':
return self
return None
def load_module(self, fullname):
if self in sys.meta_path:
sys.meta_path.remove(self)
try:
import importlib
mod = importlib.import_module(fullname)
finally:
if self not in sys.meta_path:
sys.meta_path.insert(0, self)
if _try_patch_vae_cpu():
self.done = True
return mod
class _ModelMgmtWatcher:
def __init__(self):
self.done = False
def find_module(self, fullname, path=None):
if self.done:
return None
if fullname == 'comfy.model_management':
return self
return None
def load_module(self, fullname):
if self in sys.meta_path:
sys.meta_path.remove(self)
try:
import importlib
mod = importlib.import_module(fullname)
finally:
if self not in sys.meta_path:
sys.meta_path.insert(0, self)
if _try_patch_text_encoder_device():
self.done = True
return mod
# === GPU MEMORY CLEANUP HOOK ===
# Patch model_management.load_models_gpu to clean up before loading
_load_patched = False
def _try_patch_load_models():
"""Add GPU memory cleanup before model loading."""
global _load_patched
if _load_patched:
return True
mm = sys.modules.get('comfy.model_management')
if mm is None:
return False
_original_load = getattr(mm, 'load_models_gpu', None)
if _original_load is None:
return False
def _bc250_load_models_gpu(models, *args, **kwargs):
"""Clean GPU cache before loading models to prevent memory pressure hangs."""
gc.collect()
torch.cuda.empty_cache()
return _original_load(models, *args, **kwargs)
mm.load_models_gpu = _bc250_load_models_gpu
_load_patched = True
logger.warning("[BC-250] GPU memory cleanup hook installed (load_models_gpu)")
return True
# === STARTUP PRELOAD ===
_PRELOAD_CLIP = os.environ.get("BC250_PRELOAD_CLIP", "Qwen_3_4b-Q8_0.gguf")
_PRELOAD_UNET = os.environ.get("BC250_PRELOAD_UNET", "z_image_turbo-Q5_K_S.gguf")
_PRELOAD_VAE = os.environ.get("BC250_PRELOAD_VAE", "ae.safetensors")
_PRELOAD_PORT = int(os.environ.get("BC250_PRELOAD_PORT", "8188"))
_PRELOAD_ENABLED = os.environ.get("BC250_PRELOAD", "1") == "1"
def _preload_models():
"""Background thread: wait for ComfyUI server, then submit a warmup prompt."""
import urllib.request
import urllib.error
url = f"http://127.0.0.1:{_PRELOAD_PORT}"
# Wait for server to be ready (max 120s)
logger.warning("[BC-250] Preload: waiting for ComfyUI server...")
for _ in range(240):
try:
urllib.request.urlopen(f"{url}/api/system_stats", timeout=2)
break
except (urllib.error.URLError, OSError, ConnectionRefusedError):
_time.sleep(0.5)
else:
logger.warning("[BC-250] Preload: server not ready after 120s, skipping")
return
logger.warning("[BC-250] Preload: server ready, submitting warmup prompt...")
warmup = {
"1": {"class_type": "CLIPLoaderGGUF", "inputs": {"clip_name": _PRELOAD_CLIP, "type": "lumina2"}},
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "warmup", "clip": ["1", 0]}},
"3": {"class_type": "CLIPTextEncode", "inputs": {"text": "", "clip": ["1", 0]}},
"4": {"class_type": "UnetLoaderGGUF", "inputs": {"unet_name": _PRELOAD_UNET}},
"5": {"class_type": "EmptyLatentImage", "inputs": {"width": 64, "height": 64, "batch_size": 1}},
"6": {"class_type": "KSampler", "inputs": {
"seed": 1, "steps": 1, "cfg": 1.0, "sampler_name": "euler",
"scheduler": "normal", "denoise": 1.0,
"model": ["4", 0], "positive": ["2", 0], "negative": ["3", 0], "latent_image": ["5", 0]
}},
"7": {"class_type": "VAELoader", "inputs": {"vae_name": _PRELOAD_VAE}},
"8": {"class_type": "VAEDecode", "inputs": {"samples": ["6", 0], "vae": ["7", 0]}},
"9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "_warmup", "images": ["8", 0]}}
}
payload = json.dumps({"prompt": warmup}).encode("utf-8")
req = urllib.request.Request(
f"{url}/api/prompt",
data=payload,
headers={"Content-Type": "application/json"},
method="POST"
)
try:
resp = urllib.request.urlopen(req, timeout=10)
data = json.loads(resp.read())
prompt_id = data.get("prompt_id", "unknown")
logger.warning(f"[BC-250] Preload: warmup prompt queued (id={prompt_id})")
# Wait for completion (max 5min)
for _ in range(300):
_time.sleep(1)
try:
hist_resp = urllib.request.urlopen(f"{url}/api/history/{prompt_id}", timeout=5)
hist = json.loads(hist_resp.read())
if prompt_id in hist:
logger.warning("[BC-250] Preload: all models loaded and cached. Ready for user prompts.")
return
except Exception:
pass
logger.warning("[BC-250] Preload: warmup timed out after 5min")
except Exception as e:
logger.warning(f"[BC-250] Preload: warmup failed: {e}")
def _start_preload_thread():
if not _PRELOAD_ENABLED:
logger.warning("[BC-250] Preload: disabled (BC250_PRELOAD=0)")
return
t = threading.Thread(target=_preload_models, daemon=True, name="BC250-Preload")
t.start()
logger.warning("[BC-250] Preload: background warmup thread started")
# === STARTUP PRELOAD ===
_PRELOAD_CLIP = os.environ.get("BC250_PRELOAD_CLIP", "Qwen_3_4b-Q8_0.gguf")
_PRELOAD_UNET = os.environ.get("BC250_PRELOAD_UNET", "z_image_turbo-Q5_K_S.gguf")
_PRELOAD_VAE = os.environ.get("BC250_PRELOAD_VAE", "ae.safetensors")
_PRELOAD_PORT = int(os.environ.get("BC250_PRELOAD_PORT", "8188"))
_PRELOAD_ENABLED = os.environ.get("BC250_PRELOAD", "1") == "1"
def _preload_models():
"""Background thread: wait for ComfyUI server, then submit a warmup prompt."""
import urllib.request
import urllib.error
url = f"http://127.0.0.1:{_PRELOAD_PORT}"
# Wait for server to be ready (max 120s)
logger.warning("[BC-250] Preload: waiting for ComfyUI server...")
for _ in range(240):
try:
urllib.request.urlopen(f"{url}/api/system_stats", timeout=2)
break
except (urllib.error.URLError, OSError, ConnectionRefusedError):
_time.sleep(0.5)
else:
logger.warning("[BC-250] Preload: server not ready after 120s, skipping")
return
logger.warning("[BC-250] Preload: server ready, submitting warmup prompt...")
warmup = {
"1": {"class_type": "CLIPLoaderGGUF", "inputs": {"clip_name": _PRELOAD_CLIP, "type": "lumina2"}},
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "warmup", "clip": ["1", 0]}},
"3": {"class_type": "CLIPTextEncode", "inputs": {"text": "", "clip": ["1", 0]}},
"4": {"class_type": "UnetLoaderGGUF", "inputs": {"unet_name": _PRELOAD_UNET}},
"5": {"class_type": "EmptyLatentImage", "inputs": {"width": 64, "height": 64, "batch_size": 1}},
"6": {"class_type": "KSampler", "inputs": {
"seed": 1, "steps": 1, "cfg": 1.0, "sampler_name": "euler",
"scheduler": "normal", "denoise": 1.0,
"model": ["4", 0], "positive": ["2", 0], "negative": ["3", 0], "latent_image": ["5", 0]
}},
"7": {"class_type": "VAELoader", "inputs": {"vae_name": _PRELOAD_VAE}},
"8": {"class_type": "VAEDecode", "inputs": {"samples": ["6", 0], "vae": ["7", 0]}},
"9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "_warmup", "images": ["8", 0]}}
}
payload = json.dumps({"prompt": warmup}).encode("utf-8")
req = urllib.request.Request(
f"{url}/api/prompt",
data=payload,
headers={"Content-Type": "application/json"},
method="POST"
)
try:
resp = urllib.request.urlopen(req, timeout=10)
data = json.loads(resp.read())
prompt_id = data.get("prompt_id", "unknown")
logger.warning(f"[BC-250] Preload: warmup prompt queued (id={prompt_id})")
# Wait for completion (max 5min)
for _ in range(300):
_time.sleep(1)
try:
hist_resp = urllib.request.urlopen(f"{url}/api/history/{prompt_id}", timeout=5)
hist = json.loads(hist_resp.read())
if prompt_id in hist:
logger.warning("[BC-250] Preload: all models loaded and cached. Ready for user prompts.")
return
except Exception:
pass
logger.warning("[BC-250] Preload: warmup timed out after 5min")
except Exception as e:
logger.warning(f"[BC-250] Preload: warmup failed: {e}")
def _start_preload_thread():
if not _PRELOAD_ENABLED:
logger.warning("[BC-250] Preload: disabled (BC250_PRELOAD=0)")
return
t = threading.Thread(target=_preload_models, daemon=True, name="BC250-Preload")
t.start()
logger.warning("[BC-250] Preload: background warmup thread started")
# === INSTALL ===
def _prewarm_gpu():
try:
if not torch.cuda.is_available():
return
dummy = torch.zeros(1, device='cuda')
_ = dummy + 1
torch.cuda.synchronize()
del dummy
torch.cuda.empty_cache()
logger.warning("[BC-250] GPU pre-warmed (context + allocator ready)")
except Exception as e:
logger.warning(f"[BC-250] GPU pre-warm failed: {e}")
def install():
# Mmap pre-clone patch
torch.nn.Module._apply = _bc250_safe_apply
logger.warning("[BC-250] Mmap pre-clone patch installed (XNACK workaround)")
# Softmax patches
torch.nn.functional.softmax = patched_softmax
torch.Tensor.softmax = patched_tensor_softmax
torch.nn.functional.scaled_dot_product_attention = patched_sdpa
logger.warning(f"[BC-250] Softmax monkey-patch installed (threshold={SAFE_SOFTMAX_THRESHOLD})")
# GGUF deferred cast_bias_weight patch
sys.meta_path.insert(0, _GGUFImportWatcher())
logger.warning("[BC-250] GGUF CPU-dequant hook registered (cast_bias_weight)")
# Text encoder CPU patch
sys.meta_path.insert(0, _ModelMgmtWatcher())
# VAE CPU-only patch
sys.meta_path.insert(0, _SDModuleWatcher())
# Try immediate patches if modules already loaded
_try_patch_gguf()
_try_patch_text_encoder_device()
_try_patch_vae_cpu()
_try_patch_load_models()
# Pre-warm GPU
_prewarm_gpu()
# Start background preload thread
_start_preload_thread()
install()