#!/bin/bash
# Apply transformers v5 API fix to ComfyUI-Trellis2
# Fixes: AttributeError: 'DINOv3ViTModel' object has no attribute 'layer'
# Reference: https://github.com/visualbruno/ComfyUI-Trellis2/issues/144

TRELLIS2_DIR="${1:-custom_nodes/ComfyUI-Trellis2/trellis2}"

if [ ! -d "$TRELLIS2_DIR" ]; then
    echo "ERROR: Trellis2 directory not found: $TRELLIS2_DIR"
    echo "Usage: $0 [path/to/ComfyUI-Trellis2/trellis2]"
    exit 1
fi

echo "Patching DINOv3 model.layer → model.model.layer in:"
for f in \
    modules/image_feature_extractor.py \
    trainers/flow_matching/mixins/image_conditioned.py \
    trainers/flow_matching/mixins/image_conditioned_proj.py
do
    file="$TRELLIS2_DIR/$f"
    if [ -f "$file" ]; then
        if grep -q 'self\.model\.layer' "$file"; then
            sed -i 's/self\.model\.layer/self.model.model.layer/g' "$file"
            echo "  ✓ $f"
        else
            echo "  - $f (already patched)"
        fi
    else
        echo "  ✗ $f (not found)"
    fi
done

echo "Done."