#!/bin/bash
# Add timing debug prints to load_diffusion_model_state_dict in sd.py

SD_PY="/home/fabian/ComfyUI/comfy/sd.py"

# Check if already patched  
if grep -q 'BC250_DEBUG' "$SD_PY"; then
    echo "Already has debug prints"
    exit 0
fi

# Find the model_type FLOW line and add debug before get_model call
# The sequence is roughly:
# 1. logging.info("model_type ...")  
# 2. model = model_config.get_model(new_sd, "")
# 3. model.load_model_weights(...)
# 4. model.to(offload_device) or model_patcher creation

# Add debug after "model_type" print
python3 -c "
import re

with open('$SD_PY', 'r') as f:
    content = f.read()

# Find 'model_type FLOW' or similar logging line and surrounding code
# Add timing around get_model, load_model_weights, etc.
target = 'model = model_config.get_model(new_sd, \"\")'
if target in content:
    replacement = '''import time as _t; _ts = _t.time(); logging.warning(\"[BC250_DEBUG] Creating model skeleton...\")  # BC250_DEBUG
    model = model_config.get_model(new_sd, \"\")
    logging.warning(f\"[BC250_DEBUG] Model skeleton created in {_t.time()-_ts:.1f}s\")  # BC250_DEBUG'''
    content = content.replace(target, replacement, 1)
    
target2 = 'model.load_model_weights(new_sd, prefix)'
if target2 in content:
    replacement2 = '''logging.warning(\"[BC250_DEBUG] Loading model weights...\")  # BC250_DEBUG
    _ts2 = _t.time()
    model.load_model_weights(new_sd, prefix)
    logging.warning(f\"[BC250_DEBUG] Model weights loaded in {_t.time()-_ts2:.1f}s\")  # BC250_DEBUG'''
    content = content.replace(target2, replacement2, 1)

with open('$SD_PY', 'w') as f:
    f.write(content)
print('Debug prints added')
"

echo ""
echo "=== Verify ==="
grep -n 'BC250_DEBUG' "$SD_PY"