---
# AI Toolkit config — KARATE WEINER character LoRA for FLUX.2 [klein] 4B.
# Derived from klein-build-small-starter-ref/configs/my_lora_klein_4b.yaml
# (character variant: 17 images, 1000 steps per the starter's quick-knobs).
#
# RunPod (AI Toolkit by Ostris template, RTX 4090 / L40S, >=80GB volume):
#   upload datasets/weiner -> /workspace/datasets/weiner
#   upload this file       -> /workspace/configs/weiner_klein_4b.yaml
#   cd /app/ai-toolkit && python run.py /workspace/configs/weiner_klein_4b.yaml
#
# Pick the checkpoint by SAMPLE IMAGES (expect the peak around step 500-1000),
# not by loss.
job: "extension"
config:
  name: "weiner_klein_4b_v1"
  process:
    - type: "diffusion_trainer"
      training_folder: "/app/ai-toolkit/output"
      device: "cuda"

      # Trigger word — every caption .txt starts with this exact token.
      trigger_word: "KW1ENER"

      network:
        type: "lora"
        linear: 128
        linear_alpha: 64
        conv: 64
        conv_alpha: 32

      save:
        dtype: "bf16"
        save_every: 250
        max_step_saves_to_keep: 8

      datasets:
        - folder_path: "/workspace/datasets/weiner"
          caption_ext: "txt"
          caption_dropout_rate: 0.05
          shuffle_tokens: false
          cache_latents_to_disk: true
          resolution:
            - 512
            - 768
            - 1024

      train:
        batch_size: 1
        steps: 1000                     # character LoRA, 17 images
        gradient_accumulation_steps: 1
        gradient_checkpointing: true
        noise_scheduler: "flowmatch"    # required for klein
        optimizer: "adamw8bit"
        timestep_type: "shift"
        content_or_style: "balanced"
        lr: 0.0001
        optimizer_params:
          weight_decay: 0.00015
        dtype: "bf16"

      model:
        # ALWAYS the BASE (un-distilled) checkpoint. Distilled won't fine-tune.
        name_or_path: "black-forest-labs/FLUX.2-klein-base-4B"
        # REQUIRED for klein (ai-toolkit issue #691).
        arch: "flux2_klein_4b"
        quantize: true
        low_vram: false

      sample:
        sampler: "flowmatch"
        sample_every: 250
        width: 1024
        height: 1024
        seed: 42
        walk_seed: true
        guidance_scale: 4
        sample_steps: 20
        prompts:
          - "KW1ENER. the character bows respectfully toward a pink plastic lawn flamingo in a backyard, full body"
          - "KW1ENER. close-up portrait of the character with a shocked expression, plain background"
          - "KW1ENER. a 4-panel comic strip in a 2x2 grid with thin white gutters: panel 1 the character bows to a lawn flamingo, panel 2 the character wobbles on one leg, panel 3 the character lies face down on grass, panel 4 the character balances serenely on one leg"

meta:
  name: "weiner_klein_4b"
  version: "1.0"