job: extension
config:
  name: train_changehair
  process:
  - type: sd_trainer
    training_folder: output
    device: cuda:0
    network:
      type: lora
      linear: 128
      linear_alpha: 128
    save:
      dtype: float16
      save_every: 250
      max_step_saves_to_keep: 20
      push_to_hub: true
    hf_repo_id: DVLe/ChangeHair
    hf_private: false
    datasets:
    - folder_path: /workspace/output
      control_path: /workspace/input
      caption_ext: txt
      caption_dropout_rate: 0.05
      cache_latents_to_disk: true
      resolution:
      - 512
      - 768
      - 1024
    train:
      batch_size: 1
      steps: 4000
      gradient_accumulation_steps: 1
      train_unet: true
      train_text_encoder: false
      gradient_checkpointing: true
      noise_scheduler: flowmatch
      optimizer: adamw8bit
      lr: 5.0e-05
      timestep_type: weighted
      dtype: bf16
    model:
      name_or_path: black-forest-labs/FLUX.1-Kontext-dev
      arch: flux_kontext
      quantize: true
    sample:
      sampler: flowmatch
      sample_every: 250
      width: 1328
      height: 800
      prompts:
      - Convert hairstyle to bangs updo. Keep the face, facial features, and identity
        exactly the same. Do not change skin tone, eyes, or facial structure. Only
        alter the hairstyle and make it look natural and realistic.  --ctrl_img /workspace/input/bangs_updo_1.png
      - Convert hairstyle to afro. Keep the face, facial features, and identity exactly
        the same. Do not change skin tone, eyes, or facial structure. Only alter the
        hairstyle and make it look natural and realistic.  --ctrl_img /workspace/input/afro_2.png
      - Convert hairstyle to classic bridal crown. Keep the face, facial features,
        and identity exactly the same. Do not change skin tone, eyes, or facial structure.
        Only alter the hairstyle and make it look natural and realistic.  --ctrl_img
        /workspace/input/classic_bridal_crown_3.png
      - Convert hairstyle to fishtail braid. Keep the face, facial features, and identity
        exactly the same. Do not change skin tone, eyes, or facial structure. Only
        alter the hairstyle and make it look natural and realistic.  --ctrl_img /workspace/input/fishtail_braid_4.png
      - Convert hairstyle to long boho wave. Keep the face, facial features, and identity
        exactly the same. Do not change skin tone, eyes, or facial structure. Only
        alter the hairstyle and make it look natural and realistic.  --ctrl_img /workspace/input/long_boho_wave_5.png
      - Convert hairstyle to messy bun. Keep the face, facial features, and identity
        exactly the same. Do not change skin tone, eyes, or facial structure. Only
        alter the hairstyle and make it look natural and realistic.  --ctrl_img /workspace/input/messy_bun_6.png
      - Convert hairstyle to short curly. Keep the face, facial features, and identity
        exactly the same. Do not change skin tone, eyes, or facial structure. Only
        alter the hairstyle and make it look natural and realistic.  --ctrl_img /workspace/input/short_curly_7.png
      - Convert hairstyle to aide part straight. Keep the face, facial features, and
        identity exactly the same. Do not change skin tone, eyes, or facial structure.
        Only alter the hairstyle and make it look natural and realistic.  --ctrl_img
        /workspace/input/side_part_straight_8.png
      - Convert hairstyle to sapce buns. Keep the face, facial features, and identity
        exactly the same. Do not change skin tone, eyes, or facial structure. Only
        alter the hairstyle and make it look natural and realistic.  --ctrl_img /workspace/input/space_buns_9.png
      - Convert hairstyle to two braid. Keep the face, facial features, and identity
        exactly the same. Do not change skin tone, eyes, or facial structure. Only
        alter the hairstyle and make it look natural and realistic.  --ctrl_img /workspace/input/two_braid_10.png
      neg: ''
      seed: 42
      walk_seed: true
      guidance_scale: 2.5
      sample_steps: 28
    logging:
      log_every: 2
      use_wandb: true
      project_name: changehair-training
      run_name: train_kontext
      verbose: false
meta:
  name: train_changehair
  version: '1.0'