#!/usr/bin/env python3
"""
BC-250: Test CLIP on CPU only (bypass GPU kernel compilation).
"""
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...", flush=True)
import bc250_softmax_patch
import torch
import safetensors.torch

# Check file size
clip_path = "/home/fabian/ComfyUI/models/text_encoders/gemma2_2b_lumina2.safetensors"
fsize = os.path.getsize(clip_path) / (1024*1024*1024)
print(f"[T] CLIP file: {fsize:.2f} GB", flush=True)

# Load directly to see what's in it
print(f"[T] Loading safetensors headers...", flush=True)
t0 = time.time()
with safetensors.torch.safe_open(clip_path, framework="pt", device="cpu") as f:
    keys = list(f.keys())
    print(f"[T] Keys: {len(keys)}", flush=True)
    print(f"[T] First 5 keys: {keys[:5]}", flush=True)
    
    # Check dtype and shapes of first key
    first_tensor = f.get_tensor(keys[0])
    print(f"[T] First tensor: {keys[0]} shape={first_tensor.shape} dtype={first_tensor.dtype}", flush=True)
    
    # Check total parameter count
    total_params = 0
    for k in keys:
        t = f.get_tensor(k)
        total_params += t.numel()
    print(f"[T] Total params: {total_params/1e9:.2f}B", flush=True)

dt = time.time() - t0
print(f"[T] Loaded headers in {dt:.1f}s", flush=True)

# Now try loading CLIP with ComfyUI but force CPU
print(f"\n[T] Loading CLIP through ComfyUI (on CPU)...", flush=True)
import comfy.sd
import comfy.model_management
import folder_paths

# Monkey-patch to force CPU loading for CLIP
_orig_get_torch_device = comfy.model_management.get_torch_device
_orig_text_encoder_device = comfy.model_management.text_encoder_device
_orig_text_encoder_offload = comfy.model_management.text_encoder_offload_device

# Force text encoder to CPU
comfy.model_management.text_encoder_device = lambda: torch.device("cpu")
comfy.model_management.text_encoder_offload_device = lambda: torch.device("cpu")

t1 = time.time()
try:
    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)
    
    # Test encoding
    print(f"[T] Testing text encoding on CPU...", flush=True)
    t2 = time.time()
    tokens = clip.tokenize({"g": "a photo of a cat sitting on a windowsill"})
    print(f"[T] Tokenized in {time.time()-t2:.3f}s", flush=True)
    
    t3 = time.time()
    output = clip.encode_from_tokens_scheduled(tokens)
    cond = output[0]
    dt = time.time() - t3
    print(f"[T] CLIP encoded in {dt:.1f}s", flush=True)
    print(f"[T] Output shape: {cond.shape}, dtype: {cond.dtype}", flush=True)
    print(f"\n[T] === CLIP ON CPU WORKS! ===", flush=True)
    
except Exception as e:
    print(f"[T] ERROR: {e}", flush=True)
    import traceback
    traceback.print_exc()

# Restore
comfy.model_management.text_encoder_device = _orig_text_encoder_device
comfy.model_management.text_encoder_offload_device = _orig_text_encoder_offload

print(f"[T] Total: {time.time()-t0:.1f}s", flush=True)
os._exit(0)