Uploaded sanitized BC250/ROCm Repository.
This commit is contained in:
Fabian
@@ -0,0 +1,646 @@
|
||||
"""
|
||||
BC-250 gfx1010 Comprehensive Monkey-Patch v17
|
||||
|
||||
1. BF16 KILL: gfx1010 has NO native bf16 — force f16 everywhere
|
||||
2. Softmax: manual impl for dim > threshold (VGPR overflow fix)
|
||||
3. SDPA: manual impl for large sequences
|
||||
4. GGUF: GGMLTensor.to() patched — quantized weights ALWAYS stay on CPU
|
||||
CPU dequant → f16 → GPU transfer per layer (gfx1010 GPU can't dequant)
|
||||
5. Mmap: pre-clones non-GGUF tensor data before GPU transfer
|
||||
6. GPU: pre-warms context and caching allocator
|
||||
7. CLIP: forces text encoder to CPU (memory constraint)
|
||||
8. VAE: CPU f32 decode cached in RAM (faster than GPU on this APU)
|
||||
9. Threads: all CPU cores for intra-op parallelism
|
||||
10. Rope: force rope() to CPU — gfx1010 has NO native float64
|
||||
11. Non-blocking disabled: gfx1010 without SDMA hangs on async copies
|
||||
|
||||
v14: rope() → CPU (gfx1010 has no float64 HW)
|
||||
v15: GGUF dequant → CPU (gfx1010 GPU hangs on Q5_K bitwise ops)
|
||||
v16: GGMLTensor.to() patched to keep quantized weights on CPU
|
||||
v17: cast_to → direct .to(non_blocking=False), warmup removed
|
||||
- non_blocking=True hangs on gfx1010 (no SDMA, async HIP copy broken)
|
||||
- replaced empty_like+copy_ with direct .to() for CPU→GPU transfer
|
||||
- device_supports_non_blocking → always False for this device
|
||||
|
||||
BC-250 APU / gfx1010: XNACK disabled, no GPU page fault handling.
|
||||
Place in ComfyUI root and import as first line of main.py.
|
||||
"""
|
||||
import os
|
||||
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import sys
|
||||
import gc
|
||||
import logging
|
||||
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}")
|
||||
|
||||
# Disable torch._dynamo — gfx1010 doesn't benefit, compilation overhead is massive
|
||||
try:
|
||||
torch._dynamo.config.suppress_errors = True
|
||||
logger.warning("[BC-250] torch._dynamo: TORCHDYNAMO_DISABLE=1 + suppress_errors")
|
||||
except Exception:
|
||||
logger.warning("[BC-250] torch._dynamo: TORCHDYNAMO_DISABLE=1 (env only)")
|
||||
|
||||
|
||||
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 + BF16 KILL PATCH ===
|
||||
_original_module_apply = torch.nn.Module._apply
|
||||
|
||||
def _is_ggml_tensor(t):
|
||||
"""Check if tensor is a GGMLTensor (has GGUF quantization metadata)."""
|
||||
return hasattr(t, 'tensor_type')
|
||||
|
||||
# === GGML TENSOR CPU LOCK ===
|
||||
# Patched later when GGUF module loads (_try_patch_gguf).
|
||||
# GGMLTensor.to() is monkey-patched so quantized weights NEVER leave CPU.
|
||||
# This prevents both: GPU dequant hang AND GPU→CPU transfer hang.
|
||||
|
||||
def _bc250_safe_apply(self, fn, recurse=True):
|
||||
"""Pre-clone mmap'd CPU tensor data before GPU transfer (XNACK workaround).
|
||||
Converts ALL bf16 → f16 (gfx1010 has no native bf16 — including GGML BF16).
|
||||
BF16 GGML tensors are dequantized to f32→f16, becoming regular tensors.
|
||||
Note: GGMLTensor.clone() returns self, so quantized GGUF weights are unaffected."""
|
||||
for key, param in self._parameters.items():
|
||||
if param is None:
|
||||
continue
|
||||
# Clone CPU data for XNACK workaround (skip GGML: clone() returns self)
|
||||
if param.device.type == 'cpu' and not _is_ggml_tensor(param.data):
|
||||
param.data = param.data.clone()
|
||||
# gfx1010: no native bf16. Convert ALL bf16 → f16 (including GGML BF16)
|
||||
if param.data.dtype == torch.bfloat16:
|
||||
param.data = param.data.float().half()
|
||||
for key, buf in self._buffers.items():
|
||||
if buf is None:
|
||||
continue
|
||||
if buf.device.type == 'cpu' and not _is_ggml_tensor(buf):
|
||||
buf = buf.clone()
|
||||
if buf.dtype == torch.bfloat16:
|
||||
buf = buf.float().half()
|
||||
self._buffers[key] = buf
|
||||
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 CPU DEQUANT + GGMLTensor CPU LOCK ===
|
||||
_gguf_patched = False
|
||||
|
||||
def _try_patch_gguf():
|
||||
"""Patch GGUF for BC-250:
|
||||
1. GGMLTensor.to() → keeps quantized weights on CPU (ignore device arg)
|
||||
2. cast_bias_weight → CPU dequant, then transfer f16 result to GPU
|
||||
|
||||
gfx1010 GPU cannot dequantize Q5_K (bitwise ops hang).
|
||||
And once weights are on GPU, transferring back to CPU also hangs.
|
||||
Only safe path: weights stay CPU → dequant on CPU → f16 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)
|
||||
GGMLTensor = getattr(ops_mod, 'GGMLTensor', None)
|
||||
is_quantized_fn = getattr(dequant_mod, 'is_quantized', None)
|
||||
if GGMLLayer is None or GGMLTensor is None or is_quantized_fn is None:
|
||||
return False
|
||||
|
||||
torch_compiler_disable = getattr(ops_mod, 'torch_compiler_disable', None)
|
||||
|
||||
# === PATCH 1: GGMLTensor.to() — keep quantized on CPU ===
|
||||
_original_ggml_to = GGMLTensor.to
|
||||
|
||||
def _bc250_ggml_to(self, *args, **kwargs):
|
||||
"""Intercept .to() calls: keep quantized weights on CPU.
|
||||
Only allow dtype changes, block device changes to CUDA.
|
||||
This prevents load_models_gpu from moving GGUF weights to GPU."""
|
||||
# Check if this is a quantized tensor (has tensor_type metadata)
|
||||
if hasattr(self, 'tensor_type') and self.tensor_type is not None:
|
||||
# Parse the .to() call to extract device, dtype, non_blocking
|
||||
# Common patterns from nn.Module._apply:
|
||||
# t.to(device, dtype, non_blocking) — 3 positional
|
||||
# t.to(device) — 1 positional
|
||||
# t.to(dtype) — 1 positional (dtype)
|
||||
# t.to(device=..., dtype=..., non_blocking=...) — kwargs
|
||||
parsed_device = kwargs.get('device', None)
|
||||
parsed_dtype = kwargs.get('dtype', None)
|
||||
parsed_nb = kwargs.get('non_blocking', False)
|
||||
parsed_mem_fmt = kwargs.get('memory_format', None)
|
||||
|
||||
for a in args:
|
||||
if isinstance(a, torch.device):
|
||||
parsed_device = a
|
||||
elif isinstance(a, str):
|
||||
try:
|
||||
parsed_device = torch.device(a)
|
||||
except Exception:
|
||||
pass
|
||||
elif isinstance(a, torch.dtype):
|
||||
parsed_dtype = a
|
||||
elif isinstance(a, bool):
|
||||
parsed_nb = a
|
||||
elif a is None:
|
||||
# dtype=None from Module.to() convert function
|
||||
pass
|
||||
|
||||
# Block CUDA transfer for quantized weights — stay on CPU
|
||||
if parsed_device is not None and parsed_device.type == 'cuda':
|
||||
# Reconstruct call without device, keeping dtype/non_blocking
|
||||
remap_kwargs = {}
|
||||
if parsed_dtype is not None:
|
||||
remap_kwargs['dtype'] = parsed_dtype
|
||||
if parsed_nb:
|
||||
remap_kwargs['non_blocking'] = parsed_nb
|
||||
if parsed_mem_fmt is not None:
|
||||
remap_kwargs['memory_format'] = parsed_mem_fmt
|
||||
if remap_kwargs:
|
||||
return _original_ggml_to(self, **remap_kwargs)
|
||||
return self # No-op: was just a device move
|
||||
|
||||
return _original_ggml_to(self, *args, **kwargs)
|
||||
|
||||
GGMLTensor.to = _bc250_ggml_to
|
||||
logger.warning("[BC-250] GGMLTensor.to() patched: quantized weights locked to CPU")
|
||||
|
||||
# === PATCH 2: cast_bias_weight — CPU dequant + GPU transfer ===
|
||||
_fwd_count = [0, 0.0]
|
||||
|
||||
def _bc250_cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
|
||||
"""CPU dequant → GPU transfer. Weights are guaranteed CPU (GGMLTensor.to patched).
|
||||
Dequant on CPU via get_weight(), transfer f16 to GPU with synchronous .to().
|
||||
gfx1010 without SDMA cannot do async copies — non_blocking=False always."""
|
||||
if _fwd_count[0] == 0:
|
||||
_fwd_count[1] = _time.time()
|
||||
_fwd_count[0] += 1
|
||||
|
||||
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
|
||||
|
||||
# gfx1010: never dequant to bf16
|
||||
if dtype == torch.bfloat16:
|
||||
dtype = torch.float16
|
||||
if bias_dtype == torch.bfloat16:
|
||||
bias_dtype = torch.float16
|
||||
|
||||
is_first = _fwd_count[0] <= 3 or 499 <= _fwd_count[0] <= 505
|
||||
if is_first:
|
||||
logger.warning(f"[BC-250] CBW#{_fwd_count[0]}: device={device}, w_type={type(s.weight).__name__}, w_dev={s.weight.device}, has_tt={hasattr(s.weight, 'tensor_type')}")
|
||||
|
||||
bias = None
|
||||
if s.bias is not None:
|
||||
if is_first:
|
||||
logger.warning(f"[BC-250] CBW#{_fwd_count[0]}: get_weight(bias)...")
|
||||
bias = s.get_weight(s.bias, bias_dtype)
|
||||
if is_first:
|
||||
logger.warning(f"[BC-250] CBW#{_fwd_count[0]}: bias got, type={type(bias).__name__}, dev={bias.device}, dt={bias.dtype}")
|
||||
if type(bias) is not torch.Tensor:
|
||||
bias = bias.as_subclass(torch.Tensor)
|
||||
if bias.device != device or bias.dtype != bias_dtype:
|
||||
if is_first:
|
||||
logger.warning(f"[BC-250] CBW#{_fwd_count[0]}: bias.to({device})...")
|
||||
bias = bias.to(device=device, dtype=bias_dtype, non_blocking=False)
|
||||
if is_first:
|
||||
logger.warning(f"[BC-250] CBW#{_fwd_count[0]}: bias transferred")
|
||||
|
||||
if is_first:
|
||||
logger.warning(f"[BC-250] CBW#{_fwd_count[0]}: get_weight(weight)...")
|
||||
weight = s.get_weight(s.weight, dtype)
|
||||
if is_first:
|
||||
logger.warning(f"[BC-250] CBW#{_fwd_count[0]}: weight got, type={type(weight).__name__}, shape={list(weight.shape)}, dev={weight.device}, dt={weight.dtype}")
|
||||
if type(weight) is not torch.Tensor:
|
||||
weight = weight.as_subclass(torch.Tensor)
|
||||
if weight.device != device or weight.dtype != dtype:
|
||||
if is_first:
|
||||
logger.warning(f"[BC-250] CBW#{_fwd_count[0]}: weight.to({device}, {dtype})...")
|
||||
weight = weight.to(device=device, dtype=dtype, non_blocking=False)
|
||||
if is_first:
|
||||
logger.warning(f"[BC-250] CBW#{_fwd_count[0]}: weight transferred")
|
||||
|
||||
if _fwd_count[0] % 100 == 0:
|
||||
elapsed = _time.time() - _fwd_count[1]
|
||||
logger.warning(f"[BC-250] Layer {_fwd_count[0]}, elapsed {elapsed:.1f}s")
|
||||
|
||||
# Log all CUDA transfers to find the exact hang point
|
||||
if device is not None and hasattr(device, 'type') and device.type == 'cuda':
|
||||
logger.warning(f"[BC-250] CUDA#{_fwd_count[0]}: {list(weight.shape)} {weight.dtype} done")
|
||||
|
||||
return weight, bias
|
||||
|
||||
if torch_compiler_disable is not None:
|
||||
_bc250_cast_bias_weight = torch_compiler_disable()(_bc250_cast_bias_weight)
|
||||
|
||||
GGMLLayer.cast_bias_weight = _bc250_cast_bias_weight
|
||||
_gguf_patched = True
|
||||
logger.warning("[BC-250] GGUF patched: CPU dequant + GGMLTensor CPU-locked")
|
||||
return True
|
||||
|
||||
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
|
||||
|
||||
# === MODEL MANAGEMENT PATCHES (text encoder CPU + bf16 kill + load hook) ===
|
||||
_mm_patched = False
|
||||
|
||||
def _try_patch_model_management():
|
||||
"""Patches comfy.model_management:
|
||||
- Text encoder → CPU
|
||||
- should_use_bf16 → always False
|
||||
- unet_dtype → never returns bf16
|
||||
- load_models_gpu → cleanup + post-load bf16→f16
|
||||
"""
|
||||
global _mm_patched
|
||||
if _mm_patched:
|
||||
return True
|
||||
|
||||
mm = sys.modules.get('comfy.model_management')
|
||||
if mm is None:
|
||||
return False
|
||||
|
||||
# Text encoder on CPU
|
||||
mm.text_encoder_device = lambda: torch.device("cpu")
|
||||
mm.text_encoder_offload_device = lambda: torch.device("cpu")
|
||||
logger.warning("[BC-250] Text encoder forced to CPU")
|
||||
|
||||
# Kill bf16 globally — gfx1010 has no native bf16
|
||||
mm.should_use_bf16 = lambda *a, **kw: False
|
||||
logger.warning("[BC-250] should_use_bf16 → always False (gfx1010)")
|
||||
|
||||
# Force non_blocking=False — gfx1010 without SDMA hangs on async HIP copies
|
||||
mm.device_supports_non_blocking = lambda *a, **kw: False
|
||||
logger.warning("[BC-250] device_supports_non_blocking → always False (no SDMA)")
|
||||
|
||||
_original_unet_dtype = mm.unet_dtype
|
||||
def _bc250_unet_dtype(*args, **kwargs):
|
||||
return torch.float16 # gfx1010: always f16 (2× faster than f32, no bf16 HW)
|
||||
mm.unet_dtype = _bc250_unet_dtype
|
||||
logger.warning("[BC-250] unet_dtype patched: always f16")
|
||||
|
||||
# Ensure fp16 is recognized as available
|
||||
mm.should_use_fp16 = lambda *a, **kw: True
|
||||
logger.warning("[BC-250] should_use_fp16 → always True")
|
||||
|
||||
# Load hook: cleanup + post-load bf16 → f16 conversion
|
||||
_original_load = getattr(mm, 'load_models_gpu', None)
|
||||
if _original_load is not None:
|
||||
def _bc250_load_models_gpu(models, *args, **kwargs):
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
result = _original_load(models, *args, **kwargs)
|
||||
# Post-load: convert ALL remaining bf16 params/buffers to f16
|
||||
for m in models:
|
||||
real_model = getattr(m, 'model', None)
|
||||
if real_model is None:
|
||||
continue
|
||||
converted = 0
|
||||
for p in real_model.parameters():
|
||||
if p.dtype == torch.bfloat16:
|
||||
p.data = p.data.float().half()
|
||||
converted += 1
|
||||
for name, buf in real_model.named_buffers():
|
||||
if buf is not None and buf.dtype == torch.bfloat16:
|
||||
parts = name.split('.')
|
||||
obj = real_model
|
||||
for part in parts[:-1]:
|
||||
obj = getattr(obj, part)
|
||||
setattr(obj, parts[-1], buf.float().half())
|
||||
converted += 1
|
||||
if converted > 0:
|
||||
logger.warning(f"[BC-250] Post-load: converted {converted} bf16→f16 params/buffers")
|
||||
return result
|
||||
mm.load_models_gpu = _bc250_load_models_gpu
|
||||
logger.warning("[BC-250] GPU load hook installed (cleanup + bf16 kill)")
|
||||
|
||||
_mm_patched = True
|
||||
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_model_management():
|
||||
self.done = True
|
||||
return mod
|
||||
|
||||
# === VAE CPU CACHED ===
|
||||
_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_cached_cpu(self):
|
||||
"""Keep VAE on CPU in RAM, eval mode. No GPU transfer needed.
|
||||
On BC-250 APU: GPU VAE decode is slower than CPU (24 CUs, no SDMA).
|
||||
CPU has 12 Zen2 threads and direct RAM access — faster for VAE convolutions."""
|
||||
global _vae_cached
|
||||
if not _vae_cached:
|
||||
t0 = _time.time()
|
||||
self.first_stage_model.to(device='cpu', dtype=torch.float32)
|
||||
self.first_stage_model.eval()
|
||||
_vae_cached = True
|
||||
logger.warning(f"[BC-250] VAE cached on CPU (f32) in {_time.time()-t0:.1f}s")
|
||||
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]
|
||||
|
||||
_ensure_vae_cached_cpu(self)
|
||||
|
||||
pixel_samples = None
|
||||
with torch.no_grad():
|
||||
for x in range(samples_in.shape[0]):
|
||||
sample = samples_in[x:x+1].float().cpu()
|
||||
decoded = self.first_stage_model.decode(sample, **vae_options)
|
||||
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
|
||||
|
||||
pixel_samples = pixel_samples.movedim(1, -1)
|
||||
elapsed = _time.time() - t0
|
||||
logger.warning(f"[BC-250] VAE decode (CPU f32): {elapsed:.1f}s")
|
||||
return pixel_samples
|
||||
|
||||
def _bc250_vae_encode(self, pixel_samples):
|
||||
t0 = _time.time()
|
||||
self.throw_exception_if_invalid()
|
||||
_ensure_vae_cached_cpu(self)
|
||||
with torch.no_grad():
|
||||
pixels_in = self.process_input(pixel_samples).float().cpu()
|
||||
result = self.first_stage_model.encode(pixels_in).float()
|
||||
logger.warning(f"[BC-250] VAE encode (CPU f32): {_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: CPU f32 decode/encode (cached in RAM)")
|
||||
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
|
||||
|
||||
# === ROPE CPU PATCH (gfx1010 has no float64 hardware) ===
|
||||
_rope_patched = False
|
||||
|
||||
def _try_patch_rope():
|
||||
"""Patch rope() in flux/math.py to always compute on CPU.
|
||||
gfx1010 has no native float64 — GPU float64 ops are software-emulated and hang."""
|
||||
global _rope_patched
|
||||
if _rope_patched:
|
||||
return True
|
||||
|
||||
flux_math = sys.modules.get('comfy.ldm.flux.math')
|
||||
if flux_math is None:
|
||||
return False
|
||||
|
||||
_original_rope = getattr(flux_math, 'rope', None)
|
||||
if _original_rope is None:
|
||||
return False
|
||||
|
||||
def _bc250_rope(pos, dim, theta):
|
||||
"""Compute rope on CPU (float64 not supported on gfx1010), then move result to original device."""
|
||||
assert dim % 2 == 0
|
||||
target_device = pos.device
|
||||
device = torch.device("cpu")
|
||||
scale = torch.linspace(0, (dim - 2) / dim, steps=dim // 2, dtype=torch.float64, device=device)
|
||||
omega = 1.0 / (theta ** scale)
|
||||
out = torch.einsum("...n,d->...nd", pos.to(dtype=torch.float32, device=device), omega)
|
||||
from einops import rearrange
|
||||
out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
|
||||
out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
|
||||
return out.to(dtype=torch.float32, device=target_device)
|
||||
|
||||
flux_math.rope = _bc250_rope
|
||||
_rope_patched = True
|
||||
logger.warning("[BC-250] rope() patched: CPU computation (no float64 on gfx1010)")
|
||||
return True
|
||||
|
||||
class _FluxMathWatcher:
|
||||
def __init__(self):
|
||||
self.done = False
|
||||
def find_module(self, fullname, path=None):
|
||||
if self.done:
|
||||
return None
|
||||
if fullname == 'comfy.ldm.flux.math':
|
||||
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_rope():
|
||||
self.done = True
|
||||
return mod
|
||||
|
||||
# === WARMUP REMOVED (v17) ===
|
||||
# First prompt may be slower, subsequent prompts benefit from HIP kernel caches.
|
||||
|
||||
# === INSTALL ===
|
||||
|
||||
def _prewarm_gpu():
|
||||
try:
|
||||
if not torch.cuda.is_available():
|
||||
return
|
||||
# Warm GPU context + allocator
|
||||
dummy = torch.zeros(1, device='cuda')
|
||||
_ = dummy + 1
|
||||
torch.cuda.synchronize()
|
||||
# Warm CPU→GPU copy kernel (COMGR JIT on first transfer)
|
||||
cpu_t = torch.randn(256, 256, dtype=torch.float16)
|
||||
gpu_t = cpu_t.to('cuda')
|
||||
_ = torch.matmul(gpu_t, gpu_t.T)
|
||||
torch.cuda.synchronize()
|
||||
del dummy, cpu_t, gpu_t, _
|
||||
torch.cuda.empty_cache()
|
||||
logger.warning("[BC-250] GPU pre-warmed (context + copy + matmul)")
|
||||
except Exception as e:
|
||||
logger.warning(f"[BC-250] GPU pre-warm failed: {e}")
|
||||
|
||||
|
||||
def install():
|
||||
# Mmap pre-clone + bf16 kill patch
|
||||
torch.nn.Module._apply = _bc250_safe_apply
|
||||
logger.warning("[BC-250] Mmap pre-clone + bf16→f16 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())
|
||||
sys.meta_path.insert(0, _FluxMathWatcher())
|
||||
|
||||
# Try immediate patches
|
||||
_try_patch_gguf()
|
||||
_try_patch_model_management()
|
||||
_try_patch_vae_gpu()
|
||||
_try_patch_rope()
|
||||
|
||||
_prewarm_gpu()
|
||||
logger.warning("[BC-250] v17 ready — no warmup, first prompt may be slow")
|
||||
|
||||
install()
|
||||
Reference in New Issue
Block a user