#!/bin/bash
# Test basic tensor.to(cuda) with explicit flushing
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

# Kill any remaining python processes first
pkill -9 -f "python.*main.py" 2>/dev/null
sleep 1

timeout 120 python3 -u -c "
import sys, torch, time, os
sys.stdout.flush()
print('Starting transfer test...', flush=True)
print(f'CUDA available: {torch.cuda.is_available()}', flush=True)
print(f'Device: {torch.cuda.get_device_name(0)}', flush=True)

# Test 1: Create on GPU
print('Test 1: torch.randn(10) on GPU...', flush=True)
t0 = time.time()
a = torch.randn(10, device='cuda')
torch.cuda.synchronize()
dt = time.time()-t0
print(f'  PASS in {dt:.3f}s', flush=True)

# Test 2: Small CPU to GPU
print('Test 2: small .to(cuda)...', flush=True)
b = torch.randn(10)
t0 = time.time()
c = b.to('cuda')
torch.cuda.synchronize()
dt = time.time()-t0
print(f'  PASS in {dt:.3f}s', flush=True)

# Test 3: 1MB float
print('Test 3: 1MB float .to(cuda)...', flush=True)
d = torch.randn(256*1024)
t0 = time.time()
e = d.to('cuda')
torch.cuda.synchronize()
dt = time.time()-t0
print(f'  PASS in {dt:.3f}s', flush=True)

# Test 4: 1MB uint8
print('Test 4: 1MB uint8 .to(cuda)...', flush=True)
f = torch.randint(0, 255, (1024*1024,), dtype=torch.uint8)
t0 = time.time()
g = f.to('cuda')
torch.cuda.synchronize()
dt = time.time()-t0
print(f'  PASS in {dt:.3f}s', flush=True)

# Test 5: 10MB uint8
print('Test 5: 10MB uint8 .to(cuda)...', flush=True)
h = torch.randint(0, 255, (10*1024*1024,), dtype=torch.uint8)
t0 = time.time()
i = h.to('cuda')
torch.cuda.synchronize()
dt = time.time()-t0
print(f'  PASS in {dt:.3f}s', flush=True)

# Test 6: 100MB uint8
print('Test 6: 100MB uint8 .to(cuda)...', flush=True)
j = torch.randint(0, 255, (100*1024*1024,), dtype=torch.uint8)
t0 = time.time()
k = j.to('cuda')
torch.cuda.synchronize()
dt = time.time()-t0
print(f'  PASS in {dt:.3f}s', flush=True)

# Test 7: GGMLTensor-like subclass
print('Test 7: Custom subclass .to(cuda)...', flush=True)
class FakeTensor(torch.Tensor):
    def __new__(cls, data):
        return torch.Tensor._make_subclass(cls, data)
    def to(self, *args, **kwargs):
        new = super().to(*args, **kwargs)
        return new

raw = torch.randint(0, 255, (1024*1024,), dtype=torch.uint8)
ft = FakeTensor(raw)
t0 = time.time()
ft2 = ft.to('cuda')
torch.cuda.synchronize()
dt = time.time()-t0
print(f'  PASS in {dt:.3f}s', flush=True)

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