Fabian
87 lines
2.7 KiB
Python
87 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
BC-250: Test CLIP encoding on CPU — find out if Gemma-2 2B works.
|
|
"""
|
|
import os, sys, time, signal
|
|
|
|
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")
|
|
|
|
def timeout_handler(sig, frame):
|
|
print("\n[T] === TIMEOUT HIT ===", flush=True)
|
|
os._exit(1)
|
|
signal.signal(signal.SIGALRM, timeout_handler)
|
|
|
|
print("[T] Importing...", flush=True)
|
|
import bc250_softmax_patch
|
|
import torch
|
|
import comfy.sd
|
|
import comfy.model_management
|
|
|
|
# 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")
|
|
|
|
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: {fsize:.2f} GB", flush=True)
|
|
|
|
# Load CLIP
|
|
t0 = time.time()
|
|
clip = comfy.sd.load_clip(
|
|
ckpt_paths=[clip_path],
|
|
embedding_directory=None,
|
|
clip_type=comfy.sd.CLIPType.LUMINA2,
|
|
)
|
|
print(f"[T] CLIP loaded in {time.time()-t0:.1f}s", flush=True)
|
|
|
|
# Check the clip object
|
|
print(f"[T] CLIP type: {type(clip)}", flush=True)
|
|
print(f"[T] CLIP cond_stage_model type: {type(clip.cond_stage_model)}", flush=True)
|
|
|
|
# Tokenize with just a string
|
|
text = "a photo of a cat"
|
|
print(f"[T] Tokenizing: '{text}'", flush=True)
|
|
t1 = time.time()
|
|
tokens = clip.tokenize(text)
|
|
dt = time.time() - t1
|
|
print(f"[T] Tokenized in {dt:.3f}s", flush=True)
|
|
print(f"[T] Token keys: {list(tokens.keys()) if isinstance(tokens, dict) else type(tokens)}", flush=True)
|
|
|
|
# Encode with 120s timeout
|
|
print(f"[T] Encoding (120s timeout)...", flush=True)
|
|
signal.alarm(120)
|
|
t2 = time.time()
|
|
try:
|
|
output = clip.encode_from_tokens_scheduled(tokens)
|
|
dt = time.time() - t2
|
|
signal.alarm(0)
|
|
print(f"[T] Encoded in {dt:.1f}s", flush=True)
|
|
|
|
if isinstance(output, dict):
|
|
for k, v in output.items():
|
|
if hasattr(v, 'shape'):
|
|
print(f"[T] {k}: shape={v.shape} dtype={v.dtype}", flush=True)
|
|
else:
|
|
print(f"[T] {k}: {type(v)}", flush=True)
|
|
elif isinstance(output, (list, tuple)):
|
|
for i, v in enumerate(output):
|
|
if hasattr(v, 'shape'):
|
|
print(f"[T] [{i}]: shape={v.shape} dtype={v.dtype}", flush=True)
|
|
else:
|
|
print(f"[T] [{i}]: {type(v)}", flush=True)
|
|
|
|
print(f"\n[T] === CLIP ENCODE ON CPU: SUCCESS ===", flush=True)
|
|
except Exception as e:
|
|
signal.alarm(0)
|
|
print(f"[T] ERROR: {e}", flush=True)
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
print(f"[T] Total: {time.time()-t0:.1f}s", flush=True)
|
|
os._exit(0)
|