# Training config for Flux 2 Klein Base 9B (RECOMMENDED)
# Project: paste_pose_first
# Generated: 2026-03-25 04:02:30
# Edit this file before launching training!

# === Model ===
dit = "/workspace/models/flux2-klein-base-9b/flux-2-klein-base-9b.safetensors"
vae = "/workspace/models/flux2-klein-base-9b/ae.safetensors"
text_encoder = "/workspace/models/flux2-klein-base-9b/text_encoder/model-00001-of-00004.safetensors"
model_version = "klein-base-9b"
dataset_config = "/workspace/projects/paste_pose_first/dataset.toml"

# === Attention ===
sdpa = true
flash_attn = false
xformers = false

# === Precision & Memory ===
mixed_precision = "bf16"
gradient_checkpointing = true

# === Timestep ===
timestep_sampling = "flux2_shift"
weighting_scheme = "none"

# === LoRA ===
network_module = "networks.lora_flux_2"
network_dim = 128
network_alpha = 128

# === Optimizer ===
optimizer_type = "adamw8bit"
learning_rate = 1e-4
lr_scheduler = "cosine"
gradient_accumulation_steps = 1
max_grad_norm = 1.0
lr_warmup_steps = 250

# === Duration ===
max_train_steps = 5000

# === Save ===
save_every_n_steps = 250
output_dir = "/workspace/projects/paste_pose_first/outputs/20260325_033320"
output_name = "paste_pose_first"

# === Sampling ===
sample_prompts = "/workspace/projects/paste_pose_first/sample_prompts.txt"
sample_every_n_steps = 150

# === Dataloader ===
max_data_loader_n_workers = 2
persistent_data_loader_workers = true
seed = 42