#!/bin/bash
# Patch sd.py to add model. prefix for bare Gemma2 keys
# Insert after the lm_head.weight fix (line 1302)

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

# Check if already patched
if grep -q "bare Gemma2 keys" "$SD_PY"; then
    echo "[PATCH] Already patched"
    exit 0
fi

# Verify the anchor line exists
if ! grep -q 'clip_data\[i\]\["model.lm_head.weight"\] = clip_data\[i\].pop("lm_head.weight")' "$SD_PY"; then
    echo "[PATCH] ERROR: anchor line not found"
    exit 1
fi

# Insert the fix after the lm_head.weight line
sed -i '/clip_data\[i\]\["model.lm_head.weight"\] = clip_data\[i\].pop("lm_head.weight")/a\
\        # BC-250: Add model. prefix for bare Gemma2 keys (some exports omit it)\
\        if "layers.0.post_feedforward_layernorm.weight" in clip_data[i] and "model.layers.0.post_feedforward_layernorm.weight" not in clip_data[i]:\
\            clip_data[i] = {"model." + k: v for k, v in clip_data[i].items()}' "$SD_PY"

# Verify
if grep -q "bare Gemma2 keys" "$SD_PY"; then
    echo "[PATCH] Success — Gemma2 key prefix fix applied"
    sed -n '1298,1310p' "$SD_PY"
else
    echo "[PATCH] ERROR: patch failed"
    exit 1
fi