@ -1,34 +1,23 @@
#需要环境：
### Trianing model
# 需要环境：
export HF_HUB_OFFLINE = 1
# 强制使用本地模型缓存，避免联网下载模型
### SD v2.1
# export HF_HOME="/root/autodl-tmp/huggingface_cache"
# export MODEL_PATH="stabilityai/stable-diffusion-2-1"
### SD v1.5
# export HF_HOME="/root/autodl-tmp/huggingface_cache"
# export MODEL_PATH="runwayml/stable-diffusion-v1-5"
# SD v1.5 本地路径
export MODEL_PATH = "../../static/hf_models/hub/models--runwayml--stable-diffusion-v1-5/snapshots/451f4fe16113bff5a5d2269ed5ad43b0592e9a14"
export TASKNAME = "task001"
export TYPE = "clean" #clean or perturbed
export TYPE = "perturbed" # clean or perturbed
if [ " $TYPE " = = "clean" ] ; then
export INSTANCE_DIR = " ../../static/originals/ ${ TASKNAME } "
else
export INSTANCE_DIR = " ../../static/perturbed/ ${ TASKNAME } "
fi
export DREAMBOOTH_OUTPUT_DIR = " ../../static/hf_models/fine_tuned/ ${ TYPE } / ${ TASKNAME } "
export OUTPUT_INFER_DIR = " ../../static/model_outputs/ ${ TYPE } / ${ TASKNAME } "
export CLASS_DIR = " ../../static/class/ ${ TASKNAME } "
export COORD_DIR = " ../../static/eva_res/position/ ${ TASKNAME } "
export DREAMBOOTH_OUTPUT_DIR = " ../../static/hf_models/fine_tuned/ ${ TYPE } / ${ TASKNAME } "
export OUTPUT_INFER_DIR = " ../../static/model_outputs/ ${ TYPE } / ${ TASKNAME } "
export CLASS_DIR = " ../../static/class/ ${ TASKNAME } "
export COORD_DIR = " ../../static/eva_res/position/ ${ TASKNAME } "
# ------------------------- 自动创建依赖路径 -------------------------
echo "Creating required directories..."
mkdir -p " $INSTANCE_DIR "
mkdir -p " $DREAMBOOTH_OUTPUT_DIR "
@ -36,52 +25,43 @@ mkdir -p "$OUTPUT_INFER_DIR"
mkdir -p " $CLASS_DIR "
mkdir -p " $COORD_DIR "
# ------------------------- 自动清除旧文件 -------------------------
echo " Clearing output directory: $DREAMBOOTH_OUTPUT_DIR and $OUTPUT_INFER_DIR and $COORD_DIR "
# 查找并删除目录下的所有文件和子目录（但不删除 . 或 ..）
find " $DREAMBOOTH_OUTPUT_DIR " -mindepth 1 -delete
find " $OUTPUT_INFER_DIR " -mindepth 1 -delete
find " $COORD_DIR " -mindepth 1 -delete
# ------------------------- Fine-tune DreamBooth on images -------------------------
CUDA_VISIBLE_DEVICES = 0 accelerate launch ../finetune_infras/train_db_gen_trace.py \
--pretrained_model_name_or_path= $MODEL_PATH \
--pretrained_model_name_or_path= $MODEL_PATH \
--instance_data_dir= $INSTANCE_DIR \
--class_data_dir= $CLASS_DIR \
--output_dir= $DREAMBOOTH_OUTPUT_DIR \
--validation_image_output_dir= $OUTPUT_INFER_DIR \
--with_prior_preservation \
--prior_loss_weight= 1.0 \
--instance_prompt= "a photo of sks person" \
--class_prompt= "a photo of person" \
--train_text_encoder \
--prior_loss_weight= 0.4 \
--instance_prompt= "a selfie photo of <sks> person" \
--class_prompt= "a selfie photo of person" \
--resolution= 512 \
--train_batch_size= 1 \
--gradient_accumulation_steps= 1 \
--learning_rate= 2e-6 \
--lr_scheduler= "constant" \
--lr_warmup_steps= 0 \
--num_class_images= 200 \
--max_train_steps= 1000 \
--checkpointing_steps= 500 \
--center_crop \
--learning_rate= 5e-7 \
--lr_scheduler= "constant_with_warmup" \
--lr_warmup_steps= 50 \
--num_class_images= 100 \
--max_train_steps= 800 \
--checkpointing_steps= 400 \
--mixed_precision= bf16 \
--prior_generation_precision= bf16 \
--sample_batch_size= 5 \
--validation_prompt= "a photo of sks person" \
--num_validation_images 10 \
--validation_steps 500 \
--sample_batch_size= 5 \
--validation_prompt= "a selfie photo of <sks> person, head-and-shoulders, face looking at the camera, Eiffel Tower clearly visible behind, outdoor daytime, realistic" \
--num_validation_images= 5 \
--validation_num_inference_steps= 120 \
--validation_guidance_scale= 7.0 \
--validation_image_output_dir= $OUTPUT_INFER_DIR \
--coords_save_path= $COORD_DIR \
--coords_log_interval= 10
# ------------------------- 训练后清空 CLASS_DIR -------------------------
# 注意：这会在 accelerate launch 成功结束后执行
echo " Clearing class directory: $CLASS_DIR "
# 确保目录存在，避免清理命令失败
mkdir -p " $CLASS_DIR "
# 查找并删除目录下的所有文件和子目录（但不删除 . 或 ..）
find " $CLASS_DIR " -mindepth 1 -delete
echo "Script finished."
