Fabian
125 lines
4.3 KiB
Python
125 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
BC-250 ComfyUI GGUF Integration Test
|
|
Tests the actual ComfyUI loading pipeline step by step.
|
|
"""
|
|
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("[TEST] Importing bc250_softmax_patch...", flush=True)
|
|
import bc250_softmax_patch
|
|
|
|
print("[TEST] Importing torch...", flush=True)
|
|
t0 = time.time()
|
|
import torch
|
|
print(f"[TEST] torch ready in {time.time()-t0:.1f}s", flush=True)
|
|
|
|
print(f"[TEST] CUDA available: {torch.cuda.is_available()}", flush=True)
|
|
print(f"[TEST] Device: {torch.cuda.get_device_name(0)}", flush=True)
|
|
|
|
# Step 1: Load GGUF using ComfyUI-GGUF loader
|
|
print(f"\n[TEST] === Step 1: gguf_sd_loader ===", flush=True)
|
|
|
|
# Import ComfyUI-GGUF properly as a package
|
|
import importlib
|
|
custom_nodes_path = "/home/fabian/ComfyUI/custom_nodes"
|
|
if custom_nodes_path not in sys.path:
|
|
sys.path.insert(0, custom_nodes_path)
|
|
|
|
# Force import as package
|
|
gguf_pkg = importlib.import_module("ComfyUI-GGUF")
|
|
from importlib import import_module
|
|
gguf_loader = import_module("ComfyUI-GGUF.loader")
|
|
gguf_dequant = import_module("ComfyUI-GGUF.dequant")
|
|
gguf_ops = import_module("ComfyUI-GGUF.ops")
|
|
|
|
gguf_sd_loader = gguf_loader.gguf_sd_loader
|
|
|
|
t1 = time.time()
|
|
sd, extra = gguf_sd_loader("/home/fabian/ComfyUI/models/unet/z_image_turbo-Q5_K_S.gguf")
|
|
dt = time.time() - t1
|
|
print(f"[TEST] State dict loaded in {dt:.1f}s", flush=True)
|
|
print(f"[TEST] Keys: {len(sd)}", flush=True)
|
|
print(f"[TEST] Architecture: {extra.get('arch_str')}", flush=True)
|
|
|
|
# Check some tensor info
|
|
is_quantized = gguf_dequant.is_quantized
|
|
q_count = sum(1 for v in sd.values() if is_quantized(v))
|
|
print(f"[TEST] Quantized tensors: {q_count}/{len(sd)}", flush=True)
|
|
|
|
# Step 2: Test a single dequantize on CPU
|
|
print(f"\n[TEST] === Step 2: Single tensor dequant ===", flush=True)
|
|
dequantize_tensor = gguf_dequant.dequantize_tensor
|
|
for k, v in sd.items():
|
|
if is_quantized(v):
|
|
print(f"[TEST] Dequantizing: {k} shape={v.tensor_shape} type={v.tensor_type}", flush=True)
|
|
t2 = time.time()
|
|
result = dequantize_tensor(v, dtype=torch.float16)
|
|
dt = time.time() - t2
|
|
print(f"[TEST] Done in {dt:.3f}s -> {result.shape} {result.dtype}", flush=True)
|
|
|
|
# Move to GPU
|
|
t3 = time.time()
|
|
gpu = result.to("cuda:0")
|
|
torch.cuda.synchronize()
|
|
dt2 = time.time() - t3
|
|
print(f"[TEST] GPU transfer in {dt2:.3f}s", flush=True)
|
|
del gpu, result
|
|
break
|
|
|
|
# Step 3: Test loading the model through ComfyUI model management
|
|
print(f"\n[TEST] === Step 3: ComfyUI model loading ===", flush=True)
|
|
try:
|
|
import comfy.sd
|
|
import comfy.model_management
|
|
|
|
print(f"[TEST] Loading model config...", flush=True)
|
|
t4 = time.time()
|
|
|
|
# Use the GGMLOps
|
|
GGMLOps = gguf_ops.GGMLOps
|
|
|
|
# Try to load via comfy's model loading
|
|
import comfy.supported_models
|
|
import comfy.model_patcher
|
|
|
|
# Detect model config from state dict
|
|
print(f"[TEST] Detecting model type...", flush=True)
|
|
model_config = comfy.model_detection.model_config_from_unet(sd, "")
|
|
print(f"[TEST] Model config: {type(model_config).__name__}", flush=True)
|
|
|
|
# Load into model skeleton
|
|
print(f"[TEST] Loading into model skeleton...", flush=True)
|
|
t5 = time.time()
|
|
model = model_config.get_model(sd, "", device=comfy.model_management.unet_offload_device())
|
|
model.model_config = model_config
|
|
print(f"[TEST] Model skeleton in {time.time()-t5:.1f}s", flush=True)
|
|
|
|
# Set operations
|
|
print(f"[TEST] Setting model operations...", flush=True)
|
|
ops = GGMLOps()
|
|
model.model.diffusion_model = comfy.ops.load_model_gpu(model.model.diffusion_model, ops.__class__)
|
|
|
|
# Load state dict
|
|
print(f"[TEST] Loading state dict into model...", flush=True)
|
|
t6 = time.time()
|
|
model.model.diffusion_model.load_state_dict(sd, strict=False)
|
|
dt = time.time() - t6
|
|
print(f"[TEST] State dict loaded in {dt:.1f}s", flush=True)
|
|
|
|
print(f"[TEST] Total model load: {time.time()-t4:.1f}s", flush=True)
|
|
|
|
except Exception as e:
|
|
print(f"[TEST] Error in Step 3: {type(e).__name__}: {e}", flush=True)
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
print(f"\n[TEST] COMPLETE in {time.time()-t0:.1f}s total", flush=True)
|
|
os._exit(0)
|