Fabian
100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
BC-250: Quick test - load GGUF into model via ComfyUI's actual pipeline.
|
|
Run from ComfyUI directory with venv active.
|
|
"""
|
|
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("[T] Importing patch...", flush=True)
|
|
import bc250_softmax_patch
|
|
|
|
print("[T] Importing comfy modules...", flush=True)
|
|
t0 = time.time()
|
|
|
|
import torch
|
|
import comfy.sd
|
|
import comfy.model_management
|
|
import comfy.utils
|
|
|
|
print(f"[T] Imports done in {time.time()-t0:.1f}s", flush=True)
|
|
|
|
# Load the GGUF model using ComfyUI's own loading
|
|
UNET_PATH = "/home/fabian/ComfyUI/models/unet/z_image_turbo-Q5_K_S.gguf"
|
|
|
|
print(f"\n[T] Loading GGUF unet...", flush=True)
|
|
t1 = time.time()
|
|
|
|
try:
|
|
# ComfyUI-GGUF has a hyphen, need to import differently
|
|
import importlib
|
|
gguf_nodes_path = "/home/fabian/ComfyUI/custom_nodes/ComfyUI-GGUF"
|
|
if gguf_nodes_path not in sys.path:
|
|
sys.path.insert(0, gguf_nodes_path)
|
|
|
|
# Re-check if GGUF patch got applied after import
|
|
import nodes as gguf_nodes_mod
|
|
import ops as gguf_ops_mod
|
|
import dequant as gguf_dequant_mod
|
|
|
|
# Register the dequant module so our patch can find it
|
|
sys.modules['custom_nodes.ComfyUI_GGUF.dequant'] = gguf_dequant_mod
|
|
sys.modules['custom_nodes.ComfyUI_GGUF.ops'] = gguf_ops_mod
|
|
|
|
bc250_softmax_patch._try_patch_gguf()
|
|
|
|
# Now load the model manually
|
|
from loader import gguf_sd_loader
|
|
from ops import GGMLOps, GGMLTensor
|
|
|
|
print(f"[T] Loading GGUF state dict...", flush=True)
|
|
t2 = time.time()
|
|
sd, extra = gguf_sd_loader(UNET_PATH)
|
|
print(f"[T] State dict loaded in {time.time()-t2:.1f}s ({len(sd)} keys, arch={extra.get('arch_str')})", flush=True)
|
|
|
|
# Now detect model config
|
|
import comfy.model_detection
|
|
print(f"[T] Detecting model config...", flush=True)
|
|
t3 = time.time()
|
|
|
|
# model_config_from_unet needs the state dict
|
|
parameters = comfy.utils.calculate_parameters(sd)
|
|
print(f"[T] Parameters: {parameters/1e9:.2f}B", flush=True)
|
|
|
|
unet_dtype = torch.float16
|
|
load_device = comfy.model_management.get_torch_device()
|
|
print(f"[T] Load device: {load_device}", flush=True)
|
|
|
|
model_config = comfy.model_detection.model_config_from_unet(sd, "")
|
|
print(f"[T] Model config: {type(model_config).__name__}", flush=True)
|
|
|
|
if model_config is None:
|
|
print(f"[T] ERROR: Could not detect model config!", flush=True)
|
|
else:
|
|
model_config.custom_operations = GGMLOps
|
|
|
|
print(f"[T] Creating model...", flush=True)
|
|
t4 = time.time()
|
|
model = model_config.get_model(sd, "", device=comfy.model_management.unet_offload_device())
|
|
print(f"[T] Model created in {time.time()-t4:.1f}s", flush=True)
|
|
|
|
print(f"[T] Loading state dict into model...", flush=True)
|
|
t5 = time.time()
|
|
model.load_model_weights(sd, "")
|
|
print(f"[T] Weights loaded in {time.time()-t5:.1f}s", flush=True)
|
|
|
|
print(f"\n[T] === UNET LOADED SUCCESSFULLY ===", flush=True)
|
|
|
|
except Exception as e:
|
|
print(f"[T] ERROR: {type(e).__name__}: {e}", flush=True)
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
print(f"[T] Total: {time.time()-t0:.1f}s", flush=True)
|
|
os._exit(0)
|