pipeline.load_lora_weights(
    "nbeerbower/FIXBODYr128-QwenImageEdit2509",
    weight_name="pytorch_lora_weights.safetensors"
)

# Aggressive Configuration for FIXBODY - 86 samples

# Model settings
pretrained_model_name_or_path: "Qwen/Qwen-Image-Edit-2509"

# Dataset settings
dataset_name: "nbeerbower/FIXBODY"
dataset_split: "train"
max_samples: null

# Prompt settings
use_prompt_directly: true
prompt_template: "{prompt}"

# LoRA configuration - MAX CAPACITY
lora_rank: 128  # Maximum practical rank
lora_alpha: 256  # 2x rank
lora_dropout: 0.0  # No dropout - let it fit
target_modules: ["to_k", "to_q", "to_v", "to_out.0", "ff.net.0.proj", "ff.net.2"]

# Training settings - AGGRESSIVE
train_batch_size: 1
gradient_accumulation_steps: 1  # Update every sample
num_train_epochs: 100  # Many epochs to really learn
max_train_steps: 10000
learning_rate: 1.5e-4  # Higher LR
lr_scheduler: "constant"  # No decay
lr_warmup_steps: 20

# Optimizer settings
use_8bit_adam: true
adam_beta1: 0.9
adam_beta2: 0.999
adam_weight_decay: 0.0001  # Minimal weight decay
adam_epsilon: 1e-8
max_grad_norm: 1.0

# Mixed precision
mixed_precision: "bf16"

# Checkpointing and logging
output_dir: "./output_fixbody_aggressive"
logging_dir: "logs"
checkpointing_steps: 300  # Frequent checkpoints
checkpoints_total_limit: 30  # Keep many
report_to: "wandb"
tracker_project_name: "qwen-fixbody"
run_name: "r128-aggressive-100ep"

# Data processing
save_embeddings: true
dataloader_num_workers: 1

# Memory optimization
gradient_checkpointing: true