MODEL_NAME="black-forest-labs/FLUX.2-klein-base-4B"
OUTPUT_DIR="/workspace/output/flux-klein-lora"
RESOLUTION=1024
TRAIN_BATCH_SIZE=4
GRADIENT_ACCUMULATION_STEPS=1
LEARNING_RATE=1e-4
LR_SCHEDULER="constant_with_warmup"
LR_WARMUP_STEPS=100
MAX_TRAIN_STEPS=4000
CHECKPOINTING_STEPS=500
MIXED_PRECISION="bf16"
OPTIMIZER="adamw_bf16"
LORA_RANK=16
LORA_ALPHA=16
VALIDATION_PROMPT="a professional photograph of a subject"
VALIDATION_STEPS=500