Fabian
96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
BC-250: Test CLIP on CPU only (bypass GPU kernel compilation).
|
|
"""
|
|
import os, sys, time
|
|
|
|
os.environ["HSA_OVERRIDE_GFX_VERSION"] = "10.1.0"
|
|
os.environ["HSA_ENABLE_SDMA"] = "0"
|
|
os.environ["HIP_VISIBLE_DEVICES"] = "0"
|
|
os.environ["BC250_SOFTMAX_THRESHOLD"] = "512"
|
|
|
|
sys.path.insert(0, "/home/fabian/ComfyUI")
|
|
|
|
print("[T] Importing...", flush=True)
|
|
import bc250_softmax_patch
|
|
import torch
|
|
import safetensors.torch
|
|
|
|
# Check file size
|
|
clip_path = "/home/fabian/ComfyUI/models/text_encoders/gemma2_2b_lumina2.safetensors"
|
|
fsize = os.path.getsize(clip_path) / (1024*1024*1024)
|
|
print(f"[T] CLIP file: {fsize:.2f} GB", flush=True)
|
|
|
|
# Load directly to see what's in it
|
|
print(f"[T] Loading safetensors headers...", flush=True)
|
|
t0 = time.time()
|
|
with safetensors.torch.safe_open(clip_path, framework="pt", device="cpu") as f:
|
|
keys = list(f.keys())
|
|
print(f"[T] Keys: {len(keys)}", flush=True)
|
|
print(f"[T] First 5 keys: {keys[:5]}", flush=True)
|
|
|
|
# Check dtype and shapes of first key
|
|
first_tensor = f.get_tensor(keys[0])
|
|
print(f"[T] First tensor: {keys[0]} shape={first_tensor.shape} dtype={first_tensor.dtype}", flush=True)
|
|
|
|
# Check total parameter count
|
|
total_params = 0
|
|
for k in keys:
|
|
t = f.get_tensor(k)
|
|
total_params += t.numel()
|
|
print(f"[T] Total params: {total_params/1e9:.2f}B", flush=True)
|
|
|
|
dt = time.time() - t0
|
|
print(f"[T] Loaded headers in {dt:.1f}s", flush=True)
|
|
|
|
# Now try loading CLIP with ComfyUI but force CPU
|
|
print(f"\n[T] Loading CLIP through ComfyUI (on CPU)...", flush=True)
|
|
import comfy.sd
|
|
import comfy.model_management
|
|
import folder_paths
|
|
|
|
# Monkey-patch to force CPU loading for CLIP
|
|
_orig_get_torch_device = comfy.model_management.get_torch_device
|
|
_orig_text_encoder_device = comfy.model_management.text_encoder_device
|
|
_orig_text_encoder_offload = comfy.model_management.text_encoder_offload_device
|
|
|
|
# Force text encoder to CPU
|
|
comfy.model_management.text_encoder_device = lambda: torch.device("cpu")
|
|
comfy.model_management.text_encoder_offload_device = lambda: torch.device("cpu")
|
|
|
|
t1 = time.time()
|
|
try:
|
|
clip = comfy.sd.load_clip(
|
|
ckpt_paths=[clip_path],
|
|
embedding_directory=None,
|
|
clip_type=comfy.sd.CLIPType.LUMINA2,
|
|
)
|
|
dt = time.time() - t1
|
|
print(f"[T] CLIP loaded in {dt:.1f}s", flush=True)
|
|
|
|
# Test encoding
|
|
print(f"[T] Testing text encoding on CPU...", flush=True)
|
|
t2 = time.time()
|
|
tokens = clip.tokenize({"g": "a photo of a cat sitting on a windowsill"})
|
|
print(f"[T] Tokenized in {time.time()-t2:.3f}s", flush=True)
|
|
|
|
t3 = time.time()
|
|
output = clip.encode_from_tokens_scheduled(tokens)
|
|
cond = output[0]
|
|
dt = time.time() - t3
|
|
print(f"[T] CLIP encoded in {dt:.1f}s", flush=True)
|
|
print(f"[T] Output shape: {cond.shape}, dtype: {cond.dtype}", flush=True)
|
|
print(f"\n[T] === CLIP ON CPU WORKS! ===", flush=True)
|
|
|
|
except Exception as e:
|
|
print(f"[T] ERROR: {e}", flush=True)
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
# Restore
|
|
comfy.model_management.text_encoder_device = _orig_text_encoder_device
|
|
comfy.model_management.text_encoder_offload_device = _orig_text_encoder_offload
|
|
|
|
print(f"[T] Total: {time.time()-t0:.1f}s", flush=True)
|
|
os._exit(0)
|