#!/bin/bash
# Test basic tensor.to(cuda) - does hipMemcpy work at all?
cd /home/fabian/ComfyUI
source venv/bin/activate

export HSA_OVERRIDE_GFX_VERSION=10.1.0
export HSA_ENABLE_SDMA=0
export HIP_VISIBLE_DEVICES=0
export PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.8

timeout 30 python3 -c "
import torch, time, os
print(f'CUDA available: {torch.cuda.is_available()}')
print(f'Device: {torch.cuda.get_device_name(0)}')

# Test 1: Create on GPU (already works)
print('Test 1: torch.randn on GPU...')
t0 = time.time()
a = torch.randn(10, device='cuda')
print(f'  OK in {time.time()-t0:.2f}s: {a[:3]}')

# Test 2: CPU to GPU transfer (the problematic path)
print('Test 2: tensor.to(cuda) small...')
b = torch.randn(10)
t0 = time.time()
c = b.to('cuda')
print(f'  OK in {time.time()-t0:.2f}s: {c[:3]}')

# Test 3: Larger transfer
print('Test 3: tensor.to(cuda) 1MB...')
d = torch.randn(256*1024)
t0 = time.time()
e = d.to('cuda')
print(f'  OK in {time.time()-t0:.2f}s, shape={e.shape}')

# Test 4: uint8 transfer (like GGUF)
print('Test 4: uint8 tensor.to(cuda)...')
f = torch.randint(0, 255, (1024*1024,), dtype=torch.uint8)
t0 = time.time()
g = f.to('cuda')
print(f'  OK in {time.time()-t0:.2f}s, shape={g.shape}')

# Test 5: Large uint8 (like a GGUF weight)
print('Test 5: 10MB uint8 tensor.to(cuda)...')
h = torch.randint(0, 255, (10*1024*1024,), dtype=torch.uint8)
t0 = time.time()
i = h.to('cuda')
print(f'  OK in {time.time()-t0:.2f}s, shape={i.shape}')

print('ALL TESTS PASSED - hipMemcpy works!')
os._exit(0)
" 2>&1
echo "Exit code: $?"