Ra fp16 (step 652000)

run_sdxs.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+# Продолжение sdxs-ветки на TFRecord-шардах i1 — одной командой.
+#
+#   bash run_sdxs.sh /tmp/hf/tst 10                 # 10 эпох по всем шардам в папке
+#   STEP=36000 INIT=transformer_sdxs bash run_sdxs.sh /tmp/hf/tst 10
+#   bash run_sdxs.sh /tmp/hf/tst 10 --fg            # в переднем плане (лог в консоль)
+#
+# Что уже зашито (чтобы ничего не помнить):
+#   --vae sdxs         VAE из vae_sdxs/ (асимметричный, декодер отдаёт 2x), хвоста нет,
+#                      пиксельный лосс выключен автоматически;
+#   --init             по умолчанию transformer_sdxs/ (веса sdxs-ветки в формате safetensors);
+#   --step             стартовый шаг для расписания lr (по умолчанию 5000 — как в экспорте);
+#   --batch 16 (переопределяется BATCH=8, если мало VRAM) --prefetch 256 --workers 8 --sample-every 2000 --text-len 192 --size 256;
+#   lr 1e-4 → 1e-5 по cosine за весь бюджет, без прогрева (продолжаем обученные веса);
+#   эпоха = реальное число картинок в папке (считается по заголовкам шардов).
+set -euo pipefail
+cd "$(dirname "$0")"
+
+DATA=${1:?укажи папку с TFRecord-шардами, например /tmp/hf/tst}
+EPOCHS=${2:-10}
+INIT=${INIT:-transformer_sdxs}
+STEP=${STEP:-5000}
+OUT=${OUT:-tmp/sdxs}
+BATCH=${BATCH:-16}
+LOG=${LOG:-sdxs_train.log}
+FG=0
+for a in "$@"; do [ "$a" = "--fg" ] && FG=1; done
+
+[ -e "$INIT" ] || { echo "нет стартовых весов: $INIT"; echo "скопируй папку transformer_sdxs/ (или .pt-чекпойнт и INIT=путь)"; exit 1; }
+[ -d "$DATA" ] || { echo "нет папки с шардами: $DATA"; exit 1; }
+
+CMD=(python3 src/train.py
+  --vae sdxs --init "$INIT" --step "$STEP"
+  --tfrecord "$DATA" --epochs "$EPOCHS"
+  --batch "$BATCH" --size 256 --text-len 192 --prefetch 256 --workers 8
+  --lr 1e-4 --lr-min 1e-5 --lr-warmup 0
+  --sample-every 2000 --save-every 4000 --log-every 200
+  --out "$OUT")
+
+echo "shard'ов в папке: $(ls -1 "$DATA"/*.tfrecord-* 2>/dev/null | wc -l)"
+echo "стартовые веса:   $INIT (шаг $STEP)"
+echo "бюджет:           $EPOCHS эпох"
+echo "лог:              $LOG"
+echo "команда:          ${CMD[*]}"
+echo
+
+if [ "$FG" = "1" ]; then
+  "${CMD[@]}"
+else
+  setsid nohup "${CMD[@]}" < /dev/null > "$LOG" 2>&1 &
+  sleep 20
+  echo "запущено в фоне. смотреть:  tail -f $LOG"
+fi

src/publish.py

@@ -19,7 +19,7 @@ from safetensors.torch import load_file, save_file
 ROOT = Path(__file__).resolve().parent.parent
 REPO_ID = "AiArtLab/Ra"
 
-COPY = ["README.md", "requirements.txt", ".gitattributes", ".gitignore", "docs", "src",
+COPY = ["README.md", "requirements.txt", "run_sdxs.sh", ".gitattributes", ".gitignore", "docs", "src",
         "vae/config.json", "vae/LICENSE", "vae/README.md", "vae/.gitattributes", "vae/example.py",
         "vae/ae", "text_encoder/config.json", "text_encoder/tokenizer.json",
         "text_encoder/tokenizer_config.json", "text_encoder/special_tokens_map.json",