#!/bin/bash
# Patch: Disable torch.cuda.empty_cache to prevent heap corruption on BC-250
# Add after the dynamo disable block

cd /home/fabian/ComfyUI

# 1. Add empty_cache kill after dynamo block
python3.11 -c "
import re

with open('bc250_softmax_patch.py', 'r') as f:
    content = f.read()

# Check if already patched
if 'empty_cache = lambda' in content:
    print('Already patched - empty_cache kill exists')
else:
    # Insert after the dynamo block
    marker = '    logger.warning(\"[BC-250] torch._dynamo: TORCHDYNAMO_DISABLE=1 (env only)\")'
    insert = '''

# === EMPTY_CACHE KILL ===
# BC-250 HIP allocator has heap corruption bug triggered by torch.cuda.empty_cache()
# free(): invalid next size (normal) → SIGABRT on repeated alloc/free cycles
# Disable empty_cache globally — BC-250 APU has shared memory, no eviction needed
_original_empty_cache = torch.cuda.empty_cache
torch.cuda.empty_cache = lambda: None
logger.warning(\"[BC-250] torch.cuda.empty_cache → DISABLED (HIP heap corruption fix)\")
'''
    content = content.replace(marker, marker + insert)
    
    with open('bc250_softmax_patch.py', 'w') as f:
        f.write(content)
    print('Patched: empty_cache kill added')
"

# 2. Also remove the explicit empty_cache calls in the patch itself
# (they're now no-ops anyway but cleaner to remove)
echo "Done"