#!/bin/bash
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
export AMD_LOG_LEVEL=3

echo "Quick GPU test starting..."
timeout 60 python3 -u -c "
import sys, time, os
print('Importing torch...', flush=True)
t0 = time.time()
import torch
print(f'torch imported in {time.time()-t0:.1f}s', flush=True)

print(f'CUDA: {torch.cuda.is_available()}', flush=True)
print(f'Device: {torch.cuda.get_device_name(0)}', flush=True)

# Just try allocating on GPU (no transfer)
print('Allocating on GPU...', flush=True)
t0 = time.time()
x = torch.zeros(100, device='cuda')
print(f'zeros in {time.time()-t0:.1f}s', flush=True)

# Now try a computation
print('GPU computation...', flush=True)
t0 = time.time()
y = x + 1
torch.cuda.synchronize()
print(f'add in {time.time()-t0:.1f}s, result={y[:5].tolist()}', flush=True)

# Now try the actual CPU->GPU transfer
print('CPU->GPU transfer...', flush=True)
cpu_t = torch.randn(10)
t0 = time.time()
gpu_t = cpu_t.cuda()
torch.cuda.synchronize()
print(f'transfer in {time.time()-t0:.1f}s, result={gpu_t[:3].tolist()}', flush=True)

print('DONE', flush=True)
os._exit(0)
" 2>&1
echo "Exit code: $?"