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

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)