job: extension
config:
  name: Kontext_sock_removal_v3
  process:
  - type: sd_trainer
    training_folder: output
    device: cuda:0
    trigger_word: "sock_removal"  # 触发词
    
    # 网络配置 - 采用效果好的设置
    network:
      type: lora
      linear: 16
      linear_alpha: 16
      lokr_full_rank: true      # ⭐ 从config.yaml学习
      lokr_factor: -1           # ⭐ 从config.yaml学习
      network_kwargs:
        ignore_if_contains: []
    
    # 保存设置 - 采用效果好的设置
    save:
      dtype: bf16                        # ⭐ 使用bf16
      save_every: 250
      max_step_saves_to_keep: 4          # ⭐ 保留更多检查点
      save_format: diffusers             # ⭐ 保存为diffusers格式
      push_to_hub: false
    
    # 数据集配置
    datasets:
    - folder_path: data/kontext_leg/kontext_training/images     # 目标图片（无袜）
      control_path: data/kontext_leg/kontext_training/controls  # 控制图片（有袜）
      mask_path: null
      mask_min_value: 0.1
      default_caption: ''
      caption_ext: txt
      caption_dropout_rate: 0.05
      shuffle_tokens: false
      cache_latents_to_disk: true
      resolution:               # ⭐ 包含1024分辨率
      - 512
      - 768
      - 1024
    
    # 训练参数 - 采用效果好的设置
    train:
      lr: 0.0001
      batch_size: 1
      steps: 3000                       # ⭐ 增加到3000步
      gradient_accumulation_steps: 1
      train_unet: true
      train_text_encoder: false
      gradient_checkpointing: true
      noise_scheduler: flowmatch
      optimizer: adamw8bit
      timestep_type: weighted
      content_or_style: balanced
      optimizer_params:
        weight_decay: 0.0001
      skip_first_sample: true
      unload_text_encoder: false        # ⭐ 从config.yaml学习
      linear_timesteps: true            # ⭐ 从config.yaml学习
      ema_config:                       # ⭐ 添加EMA配置
        use_ema: true
        ema_decay: 0.99
      dtype: bf16
      diff_output_preservation: false   # ⭐ 从config.yaml学习
      diff_output_preservation_multiplier: 1
      diff_output_preservation_class: person
      seed: 42
    
    # 模型配置 - 采用官方模型路径
    model:
      name_or_path: black-forest-labs/FLUX.1-Kontext-dev  # ⭐ 使用官方路径
      arch: flux_kontext
      is_flux: true
      quantize: true
      quantize_te: true
      model_kwargs: {}
      low_vram: true                    # ⭐ 低显存优化
    
    # 采样配置
    sample:
      sampler: flowmatch
      sample_every: 250
      width: 1024
      height: 1024
      prompts:
      # ⭐ 使用正确的----ctrl_img格式
      - Remove the socks from this woman ----ctrl_img data/kontext_leg/kontext_training/controls/sock_removal_001.jpg
      - Remove stockings from the legs ----ctrl_img data/kontext_leg/kontext_training/controls/sock_removal_010.jpg  
      - Remove hosiery from this person ----ctrl_img data/kontext_leg/kontext_training/controls/sock_removal_020.jpg
      neg: 'stockings, socks, pantyhose, tights, hosiery, bad anatomy, distorted'
      seed: 42
      walk_seed: true
      guidance_scale: 4
      sample_steps: 25
      num_frames: 1
      fps: 1

meta:
  name: Kontext_sock_removal_v3
  version: '1.0'