Uploaded sanitized BC250/ROCm Repository.
This commit is contained in:
Fabian
@@ -0,0 +1,539 @@
|
||||
"""
|
||||
BC-250 gfx1010 Comprehensive Monkey-Patch v11
|
||||
1. Softmax: manual impl for dim > threshold (VGPR overflow fix)
|
||||
2. SDPA: manual impl for large sequences
|
||||
3. GGUF: GPU dequant with weight cache (eliminates per-step dequant)
|
||||
4. Mmap: pre-clones non-GGUF tensor data before GPU transfer
|
||||
5. GPU: pre-warms context and caching allocator
|
||||
6. CLIP: forces text encoder to CPU (memory constraint)
|
||||
7. VAE: GPU fp16 decode with persistent caching (shared memory APU)
|
||||
8. Threads: all CPU cores for intra-op parallelism
|
||||
9. Preload: background warmup prompt on server start
|
||||
|
||||
NOTE: mlockall REMOVED — on APU with shared memory, pinning 10GB of mmap'd
|
||||
GGUF files leaves no room for GPU GTT allocations → OOM kill.
|
||||
The kernel page cache handles this correctly without mlockall.
|
||||
|
||||
v11 changes vs v10:
|
||||
- GPU dequant instead of CPU (GGUF dequant ops are pure PyTorch, run on GPU)
|
||||
- Weight cache: dequanted fp16 weights cached per-layer, reused across steps
|
||||
- mlockall() to pin process memory in RAM (no zram/swap penalty)
|
||||
- Removed duplicate preload section
|
||||
- Clean rewrite
|
||||
|
||||
BC-250 APU / gfx1010: XNACK disabled, no GPU page fault handling.
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# === THREAD CONFIGURATION ===
|
||||
_NUM_THREADS = int(os.environ.get("BC250_NUM_THREADS", str(os.cpu_count() or 12)))
|
||||
torch.set_num_threads(_NUM_THREADS)
|
||||
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 (XNACK workaround).
|
||||
Note: GGMLTensor.clone() returns self, so GGUF weights are unaffected.
|
||||
They're handled by GGMLTensor.to() which preserves metadata."""
|
||||
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: shape={list(input.shape)}, dim={dim}")
|
||||
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: shape={list(self.shape)}, dim={dim}")
|
||||
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: Q={list(query.shape)}, S={S}")
|
||||
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 WEIGHT CACHE + GPU DEQUANT ===
|
||||
_gguf_patched = False
|
||||
_weight_cache = {}
|
||||
_weight_cache_bytes = 0
|
||||
_WEIGHT_CACHE_MB = int(os.environ.get("BC250_WEIGHT_CACHE_MB", "0"))
|
||||
|
||||
def _try_patch_gguf():
|
||||
"""Patch GGMLLayer.cast_bias_weight: GPU dequant + weight caching."""
|
||||
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
|
||||
|
||||
_original_cast = getattr(GGMLLayer, 'cast_bias_weight', None)
|
||||
cache_budget = _WEIGHT_CACHE_MB * 1024 * 1024
|
||||
|
||||
def _bc250_cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
|
||||
"""GPU dequant with optional weight caching.
|
||||
|
||||
With --highvram, GGUF weights are already on GPU. Dequant happens
|
||||
via PyTorch tensor ops on GPU (parallel) instead of CPU (sequential).
|
||||
If weight cache is enabled (BC250_WEIGHT_CACHE_MB > 0), dequanted
|
||||
weights are cached per-layer to eliminate dequant on steps 2+.
|
||||
"""
|
||||
global _weight_cache_bytes
|
||||
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)
|
||||
|
||||
# Check weight cache
|
||||
if cache_budget > 0:
|
||||
cache_key = id(s)
|
||||
cached = _weight_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
# Bias
|
||||
bias = None
|
||||
if s.bias is not None:
|
||||
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)
|
||||
|
||||
# Weight: .to(device) moves GGMLTensor to GPU, get_weight dequants on GPU
|
||||
weight = s.get_weight(s.weight.to(device), dtype)
|
||||
weight = comfy.ops.cast_to(weight, dtype, device, non_blocking=non_blocking, copy=False)
|
||||
|
||||
# Cache if within budget
|
||||
if cache_budget > 0:
|
||||
entry_bytes = weight.nelement() * weight.element_size()
|
||||
if bias is not None:
|
||||
entry_bytes += bias.nelement() * bias.element_size()
|
||||
if _weight_cache_bytes + entry_bytes <= cache_budget:
|
||||
_weight_cache[cache_key] = (weight, bias)
|
||||
_weight_cache_bytes += entry_bytes
|
||||
|
||||
return weight, bias
|
||||
|
||||
GGMLLayer.cast_bias_weight = _bc250_cast_bias_weight
|
||||
|
||||
_gguf_patched = True
|
||||
cache_str = f", weight cache={_WEIGHT_CACHE_MB}MB" if cache_budget > 0 else ""
|
||||
logger.warning(f"[BC-250] GGUF patched: GPU dequant{cache_str}")
|
||||
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")
|
||||
return True
|
||||
|
||||
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
|
||||
|
||||
# === VAE GPU FP16 WITH CACHING ===
|
||||
# BC-250 = APU with shared memory. GPU VRAM = CPU RAM = same physical pool.
|
||||
# No OOM risk from "using VRAM" — it's all the same 16GB.
|
||||
# GPU fp16 VAE is ~10x faster than CPU float32.
|
||||
_vae_patched = False
|
||||
_vae_cached = False
|
||||
|
||||
def _try_patch_vae_gpu():
|
||||
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
|
||||
|
||||
def _ensure_vae_on_gpu_f16(self):
|
||||
"""Move VAE to GPU fp16 once, keep it cached. Shared memory = no OOM risk."""
|
||||
global _vae_cached
|
||||
gpu = torch.device("cuda")
|
||||
try:
|
||||
p = next(self.first_stage_model.parameters())
|
||||
already_ready = _vae_cached and p.device.type == 'cuda' and p.dtype == torch.float16
|
||||
except StopIteration:
|
||||
already_ready = False
|
||||
if not already_ready:
|
||||
logger.warning("[BC-250] Loading VAE to GPU fp16 (shared memory, will stay cached)")
|
||||
# Bypass _bc250_safe_apply (mmap pre-clone) — VAE is safetensors, not GGUF
|
||||
old_apply = torch.nn.Module._apply
|
||||
torch.nn.Module._apply = _original_module_apply
|
||||
try:
|
||||
self.first_stage_model.half().cuda()
|
||||
finally:
|
||||
torch.nn.Module._apply = old_apply
|
||||
self.first_stage_model.eval()
|
||||
_vae_cached = True
|
||||
self.disable_offload = True
|
||||
|
||||
def _bc250_vae_decode(self, samples_in, vae_options={}):
|
||||
t0 = _time.time()
|
||||
self.throw_exception_if_invalid()
|
||||
|
||||
if self.latent_dim == 2 and samples_in.ndim == 5:
|
||||
samples_in = samples_in[:, :, 0]
|
||||
|
||||
# Free GPU memory from UNET before loading VAE
|
||||
mm = sys.modules.get('comfy.model_management')
|
||||
if mm:
|
||||
mm.unload_all_models()
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
_ensure_vae_on_gpu_f16(self)
|
||||
|
||||
pixel_samples = None
|
||||
with torch.no_grad():
|
||||
for x in range(samples_in.shape[0]):
|
||||
sample = samples_in[x:x+1].to(torch.float16).cuda()
|
||||
decoded = self.first_stage_model.decode(sample, **vae_options)
|
||||
if decoded.ndim == 5:
|
||||
decoded = decoded[:, :, 0]
|
||||
out = self.process_output(decoded.float().cpu())
|
||||
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
|
||||
|
||||
pixel_samples = pixel_samples.movedim(1, -1)
|
||||
elapsed = _time.time() - t0
|
||||
logger.warning(f"[BC-250] VAE decode (GPU fp16): {elapsed:.1f}s")
|
||||
return pixel_samples
|
||||
|
||||
def _bc250_vae_encode(self, pixel_samples):
|
||||
t0 = _time.time()
|
||||
self.throw_exception_if_invalid()
|
||||
mm = sys.modules.get('comfy.model_management')
|
||||
if mm:
|
||||
mm.unload_all_models()
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
_ensure_vae_on_gpu_f16(self)
|
||||
with torch.no_grad():
|
||||
pixels_in = self.process_input(pixel_samples).to(torch.float16).cuda()
|
||||
result = self.first_stage_model.encode(pixels_in).float().cpu()
|
||||
logger.warning(f"[BC-250] VAE encode (GPU fp16): {_time.time() - t0:.1f}s")
|
||||
return result
|
||||
|
||||
VAE.decode = _bc250_vae_decode
|
||||
VAE.encode = _bc250_vae_encode
|
||||
_vae_patched = True
|
||||
logger.warning("[BC-250] VAE patched: GPU fp16 (shared memory = zero OOM risk)")
|
||||
return True
|
||||
|
||||
class _SDModuleWatcher:
|
||||
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_gpu():
|
||||
self.done = True
|
||||
return mod
|
||||
|
||||
# === GPU MEMORY CLEANUP HOOK ===
|
||||
_load_patched = False
|
||||
|
||||
def _try_patch_load_models():
|
||||
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):
|
||||
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")
|
||||
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():
|
||||
import urllib.request
|
||||
import urllib.error
|
||||
|
||||
url = f"http://127.0.0.1:{_PRELOAD_PORT}"
|
||||
|
||||
logger.warning("[BC-250] Preload: waiting for 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, skip")
|
||||
return
|
||||
|
||||
logger.warning("[BC-250] Preload: server ready, submitting warmup...")
|
||||
|
||||
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 queued (id={prompt_id})")
|
||||
|
||||
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 cached. Ready.")
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger.warning("[BC-250] Preload: warmup timed out (5min)")
|
||||
except Exception as e:
|
||||
logger.warning(f"[BC-250] Preload 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 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")
|
||||
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")
|
||||
|
||||
# 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 patch installed (threshold={SAFE_SOFTMAX_THRESHOLD})")
|
||||
|
||||
# Deferred patches via import hooks
|
||||
sys.meta_path.insert(0, _GGUFImportWatcher())
|
||||
sys.meta_path.insert(0, _ModelMgmtWatcher())
|
||||
sys.meta_path.insert(0, _SDModuleWatcher())
|
||||
|
||||
# Try immediate patches
|
||||
_try_patch_gguf()
|
||||
_try_patch_text_encoder_device()
|
||||
_try_patch_vae_gpu()
|
||||
_try_patch_load_models()
|
||||
|
||||
_prewarm_gpu()
|
||||
_start_preload_thread()
|
||||
|
||||
install()
|
||||
Reference in New Issue
Block a user