stable-diffusion-paperspace/other/fix-diffusers-root-director...

42 lines
1.1 KiB
Bash

#!/bin/bash
# Diffusers creates a pickle file with the file name as the root directory.
# This trips the WebUI's pickle checker which expects the root directory to be
# named "archive". This script renames the root directory to "archive".
if [ -z "$1" ]; then echo -e "Usage:\n./fix-diffusers-root-directory.sh [file or directory of files]" && exit 1; fi
TARGET_PATH="${@%/}"
test_zip () {
if file "$1" | grep "Zip archive" > /dev/null 2>&1; then
return 0
else
return 1
fi
}
do_convert () {
TARGET_FILE="$1"
FILENAME=$(basename "$TARGET_FILE")
NAME="${FILENAME%.*}"
EXT="${FILENAME##*.}"
NEW_NAME="${NAME}_fixed.$EXT"
echo "$TARGET_FILE -> $NEW_NAME"
# Passing the file back through torch seems to fix it
python3 -c "import torch; torch.save(torch.load('$TARGET_FILE', map_location='cpu'), '$NEW_NAME')"
}
if [ -d "$TARGET_PATH" ]; then
for file in "$TARGET_PATH"/*; do
if [ -f "$file" ] && test_zip "$file"; then
do_convert "$file"
fi
done
else
for file in $TARGET_PATH; do
if test_zip "$file"; then
do_convert "$file"
fi
done
fi