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
ROCm-Research-Archive/Scripts and Tests/component_test.py
T
2026-08-20 00:45:43 +02:00

144 lines
4.7 KiB
Python

#!/usr/bin/env python3
"""
BC-250: Test each ComfyUI component in isolation to find which one hangs.
Run from /home/fabian/ComfyUI with venv active.
"""
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")
# Timeout handler
def timeout_handler(signum, frame):
print(f"\n[TIMEOUT] Operation exceeded time limit!", flush=True)
os._exit(1)
print("[T] Importing patch...", flush=True)
import bc250_softmax_patch
print("[T] Importing comfy...", flush=True)
t0 = time.time()
import torch
import comfy.sd
import comfy.model_management
import comfy.utils
import comfy.clip_model
import folder_paths
print(f"[T] Imports done in {time.time()-t0:.1f}s", flush=True)
print(f"[T] CUDA: {torch.cuda.is_available()}, Device: {torch.cuda.get_device_name(0)}", flush=True)
# === TEST 1: Load CLIP model ===
print(f"\n{'='*60}", flush=True)
print(f"[T] TEST 1: Load CLIP model (gemma2_2b_lumina2)", flush=True)
signal.alarm(60) # 60s timeout
t1 = time.time()
try:
clip_path = os.path.join(folder_paths.get_folder_paths("clip")[0], "gemma2_2b_lumina2.safetensors")
if not os.path.exists(clip_path):
# Try text_encoders folder
for p in folder_paths.get_folder_paths("text_encoders"):
cp = os.path.join(p, "gemma2_2b_lumina2.safetensors")
if os.path.exists(cp):
clip_path = cp
break
print(f"[T] CLIP path: {clip_path}", flush=True)
print(f"[T] Loading CLIP...", flush=True)
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)
print(f"[T] CLIP type: {type(clip).__name__}", flush=True)
# Check GPU memory after CLIP load
print(f"[T] GPU VRAM after CLIP load:", flush=True)
print(f"[T] allocated: {torch.cuda.memory_allocated()/1e6:.1f} MB", flush=True)
print(f"[T] reserved: {torch.cuda.memory_reserved()/1e6:.1f} MB", flush=True)
except Exception as e:
print(f"[T] TEST 1 ERROR: {e}", flush=True)
import traceback
traceback.print_exc()
clip = None
signal.alarm(0)
# === TEST 2: Run CLIP text encoding ===
if clip is not None:
print(f"\n{'='*60}", flush=True)
print(f"[T] TEST 2: CLIP text encoding", flush=True)
signal.alarm(120) # 120s timeout
t2 = time.time()
try:
print(f"[T] Encoding: 'a cat'...", flush=True)
tokens = clip.tokenize({"g": "a cat"})
print(f"[T] Tokenized in {time.time()-t2:.3f}s", flush=True)
t2b = time.time()
print(f"[T] Running CLIP encode (this is the suspected hang point)...", flush=True)
output = clip.encode_from_tokens_scheduled(tokens)
cond, pooled = output[:2]
dt = time.time() - t2b
print(f"[T] CLIP encoded in {dt:.1f}s", flush=True)
print(f"[T] Cond shape: {cond.shape}, dtype: {cond.dtype}", flush=True)
except Exception as e:
print(f"[T] TEST 2 ERROR: {e}", flush=True)
import traceback
traceback.print_exc()
signal.alarm(0)
# === TEST 3: Load GGUF UNet ===
print(f"\n{'='*60}", flush=True)
print(f"[T] TEST 3: Load GGUF UNet", flush=True)
signal.alarm(60)
t3 = time.time()
try:
# Ensure GGUF patch is applied
bc250_softmax_patch._try_patch_gguf()
# Import GGUF nodes
gguf_path = "/home/fabian/ComfyUI/custom_nodes/ComfyUI-GGUF"
sys.path.insert(0, gguf_path)
# Use the loader directly
from loader import gguf_sd_loader
from ops import GGMLOps
unet_path = "/home/fabian/ComfyUI/models/unet/z_image_turbo-Q5_K_S.gguf"
print(f"[T] Loading GGUF state dict...", flush=True)
sd, extra = gguf_sd_loader(unet_path)
print(f"[T] State dict: {len(sd)} keys, arch={extra.get('arch_str')}", flush=True)
# Now load through comfy
print(f"[T] Creating diffusion model...", flush=True)
ops = GGMLOps()
model = comfy.sd.load_diffusion_model_state_dict(
sd, model_options={"custom_operations": ops},
metadata=extra.get("metadata", {}),
)
dt = time.time() - t3
print(f"[T] UNet loaded in {dt:.1f}s", flush=True)
if model is not None:
print(f"[T] Model type: {type(model).__name__}", flush=True)
else:
print(f"[T] WARNING: model is None!", flush=True)
except Exception as e:
print(f"[T] TEST 3 ERROR: {e}", flush=True)
import traceback
traceback.print_exc()
signal.alarm(0)
print(f"\n{'='*60}", flush=True)
print(f"[T] ALL TESTS COMPLETE in {time.time()-t0:.1f}s", flush=True)
os._exit(0)