File size: 3,182 Bytes
5f99e05 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | #!/bin/bash
# Run this on the RunPod pod terminal (after you've created the pod yourself).
# Assumes an RTX 4090 (24GB) or similar, Ubuntu + CUDA image, network volume at /workspace.
# Network volume should be at least 100GB this time -- kyrael's pod hit disk-full twice at 50GB.
set -e
cd /workspace
echo "=== 1. Clone musubi-tuner ==="
git clone https://github.com/kohya-ss/musubi-tuner
cd musubi-tuner
pip install -e .
pip install transformers accelerate qwen-vl-utils "huggingface_hub[cli]"
echo "=== 2. Log into HuggingFace ==="
echo "Paste your token when prompted -- do NOT put it directly on the command line."
hf auth login
echo "=== 3. Download the RAW Krea2 model (~24.5GB, gated -- must have accepted access on huggingface.co/krea/Krea-2-Raw) ==="
mkdir -p /workspace/models
hf download krea/Krea-2-Raw raw.safetensors --local-dir /workspace/models/krea2_raw
echo "=== 4. Download VAE + text encoder ==="
hf download Comfy-Org/Qwen-Image_ComfyUI split_files/vae/qwen_image_vae.safetensors --local-dir /workspace/models/vae
hf download Comfy-Org/Qwen3-VL text_encoders/qwen3vl_4b_bf16.safetensors --local-dir /workspace/models/text_encoder
echo "=== 5. Download vaelith dataset ==="
mkdir -p /workspace/dataset
hf download JBARU/vaelith-dataset --repo-type dataset --local-dir /workspace/dataset/vaelith
echo "=== 6. Caption dataset (Qwen2.5-VL-7B) ==="
python /workspace/caption_dataset.py /workspace/dataset/vaelith vaelith
echo "=== 6b. Clean up captioning model cache (~16GB) -- this is what caused the disk-full crashes on kyrael's run ==="
rm -rf /workspace/.cache
df -h /workspace
VAE=/workspace/models/vae/split_files/vae/qwen_image_vae.safetensors
TE=/workspace/models/text_encoder/text_encoders/qwen3vl_4b_bf16.safetensors
DIT=/workspace/models/krea2_raw/raw.safetensors
echo "=== 7. Pre-cache latents + text encoder outputs (vaelith) ==="
python src/musubi_tuner/krea2_cache_latents.py --dataset_config /workspace/dataset_vaelith.toml --vae "$VAE"
python src/musubi_tuner/krea2_cache_text_encoder_outputs.py --dataset_config /workspace/dataset_vaelith.toml --text_encoder "$TE" --batch_size 1
echo "=== 8. Train vaelith LoRA ==="
echo "Using num_repeats=3 this time (kyrael used 10, which caused overtraining/rigidity)."
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
accelerate launch --num_cpu_threads_per_process 1 --mixed_precision bf16 \
src/musubi_tuner/krea2_train_network.py \
--dit "$DIT" --vae "$VAE" \
--dataset_config /workspace/dataset_vaelith.toml \
--sdpa --mixed_precision bf16 --fp8_base --fp8_scaled \
--timestep_sampling shift --weighting_scheme none --discrete_flow_shift 2.5 \
--optimizer_type adamw8bit --learning_rate 1e-4 --gradient_checkpointing \
--max_data_loader_n_workers 2 --persistent_data_loader_workers \
--network_module networks.lora_krea2 --network_dim 32 --network_alpha 16 \
--max_train_epochs 16 --save_every_n_epochs 2 --seed 42 \
--output_dir /workspace/output/vaelith --output_name vaelith_lora
echo "=== Done. LoRA is in /workspace/output/vaelith ==="
echo "Back it up immediately with: hf upload <your-username>/vaelith-lora /workspace/output/vaelith --repo-type model --private"
|