141 lines
5.2 KiB
Python
141 lines
5.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
BC-250 GGUF Load Test - isolated from ComfyUI
|
|
Loads z_image_turbo GGUF and times each 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"
|
|
|
|
print("[TEST] Starting GGUF load test...", flush=True)
|
|
|
|
t0 = time.time()
|
|
import torch
|
|
print(f"[TEST] torch imported in {time.time()-t0:.1f}s", flush=True)
|
|
|
|
# Add ComfyUI paths
|
|
sys.path.insert(0, "/home/fabian/ComfyUI")
|
|
sys.path.insert(0, "/home/fabian/ComfyUI/custom_nodes/ComfyUI-GGUF")
|
|
|
|
t1 = time.time()
|
|
import gguf
|
|
print(f"[TEST] gguf imported in {time.time()-t1:.1f}s", flush=True)
|
|
|
|
GGUF_PATH = "/home/fabian/ComfyUI/models/unet/z_image_turbo-Q5_K_S.gguf"
|
|
|
|
# Step 1: Open and read the file
|
|
print(f"\n[TEST] Step 1: Reading GGUF file...", flush=True)
|
|
t2 = time.time()
|
|
reader = gguf.GGUFReader(GGUF_PATH)
|
|
print(f"[TEST] Reader created in {time.time()-t2:.1f}s", flush=True)
|
|
print(f"[TEST] Tensors: {len(reader.tensors)}", flush=True)
|
|
|
|
# Check architecture
|
|
arch = None
|
|
field = reader.get_field("general.architecture")
|
|
if field:
|
|
arch = str(field.parts[field.data[-1]], "utf-8")
|
|
print(f"[TEST] Architecture: {arch}", flush=True)
|
|
|
|
# Step 2: List first few tensors with types
|
|
print(f"\n[TEST] Step 2: Tensor info (first 10)...", flush=True)
|
|
qtype_counts = {}
|
|
for i, tensor in enumerate(reader.tensors):
|
|
ttype = tensor.tensor_type
|
|
tname = getattr(ttype, 'name', repr(ttype))
|
|
qtype_counts[tname] = qtype_counts.get(tname, 0) + 1
|
|
if i < 10:
|
|
print(f" {tensor.name}: type={tname} shape={list(reversed(tensor.shape))} size={tensor.data.nbytes}", flush=True)
|
|
|
|
print(f"\n[TEST] Quant type distribution:", flush=True)
|
|
for k, v in sorted(qtype_counts.items()):
|
|
print(f" {k}: {v}", flush=True)
|
|
|
|
# Step 3: Try loading one tensor to CPU
|
|
print(f"\n[TEST] Step 3: Load first Q5_K tensor to CPU...", flush=True)
|
|
import warnings
|
|
from dequant import dequantize_tensor, is_quantized
|
|
|
|
for tensor in reader.tensors:
|
|
if tensor.tensor_type == gguf.GGMLQuantizationType.Q5_K:
|
|
name = tensor.name
|
|
print(f" Tensor: {name}", flush=True)
|
|
print(f" Raw shape: {list(reversed(tensor.shape))}", flush=True)
|
|
print(f" Raw bytes: {tensor.data.nbytes}", flush=True)
|
|
|
|
t3 = time.time()
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", message="The given NumPy array is not writable")
|
|
torch_tensor = torch.from_numpy(tensor.data)
|
|
|
|
# Try to get orig shape
|
|
field_key = f"comfy.gguf.orig_shape.{name}"
|
|
field = reader.get_field(field_key)
|
|
if field:
|
|
shape = torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data))
|
|
else:
|
|
shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
|
|
|
|
print(f" Original shape: {shape}", flush=True)
|
|
print(f" Tensor dtype: {torch_tensor.dtype}, device: {torch_tensor.device}", flush=True)
|
|
|
|
# Create a GGMLTensor-like wrapper
|
|
torch_tensor.tensor_type = tensor.tensor_type
|
|
torch_tensor.tensor_shape = shape
|
|
torch_tensor.patches = []
|
|
|
|
# Dequantize on CPU
|
|
print(f" Dequantizing on CPU...", flush=True)
|
|
t4 = time.time()
|
|
result = dequantize_tensor(torch_tensor, dtype=torch.float16)
|
|
dt = time.time() - t4
|
|
print(f" Dequantized in {dt:.3f}s -> shape={result.shape} dtype={result.dtype}", flush=True)
|
|
|
|
# Try moving to GPU
|
|
print(f" Moving to GPU...", flush=True)
|
|
t5 = time.time()
|
|
gpu_result = result.to("cuda:0")
|
|
torch.cuda.synchronize()
|
|
dt2 = time.time() - t5
|
|
print(f" GPU transfer in {dt2:.3f}s", flush=True)
|
|
print(f" GPU tensor: shape={gpu_result.shape} dtype={gpu_result.dtype} device={gpu_result.device}", flush=True)
|
|
|
|
del gpu_result, result, torch_tensor
|
|
torch.cuda.empty_cache()
|
|
break
|
|
|
|
# Step 4: Try dequantizing ALL tensors on CPU
|
|
print(f"\n[TEST] Step 4: Dequantize all Q5_K tensors on CPU...", flush=True)
|
|
t_all = time.time()
|
|
count = 0
|
|
for tensor in reader.tensors:
|
|
if tensor.tensor_type == gguf.GGMLQuantizationType.Q5_K:
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", message="The given NumPy array is not writable")
|
|
tt = torch.from_numpy(tensor.data)
|
|
|
|
field_key = f"comfy.gguf.orig_shape.{tensor.name}"
|
|
field = reader.get_field(field_key)
|
|
if field:
|
|
shape = torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data))
|
|
else:
|
|
shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
|
|
|
|
tt.tensor_type = tensor.tensor_type
|
|
tt.tensor_shape = shape
|
|
tt.patches = []
|
|
|
|
result = dequantize_tensor(tt, dtype=torch.float16)
|
|
count += 1
|
|
if count % 20 == 0:
|
|
print(f" Dequantized {count} tensors ({time.time()-t_all:.1f}s)...", flush=True)
|
|
|
|
dt_all = time.time() - t_all
|
|
print(f"\n[TEST] Dequantized {count} Q5_K tensors in {dt_all:.1f}s total", flush=True)
|
|
print(f"[TEST] Average: {dt_all/max(count,1):.3f}s per tensor", flush=True)
|
|
|
|
print(f"\n[TEST] ALL DONE in {time.time()-t0:.1f}s", flush=True)
|
|
os._exit(0)
|