---
# NumZoo style LoRA — FLUX.2-klein-4B, trained on Modal via ai-toolkit.
# Trained on the BASE model (FLUX.2-klein-base-4B); the resulting LoRA is loaded
# on the DISTILLED FLUX.2-klein-4B at inference in the app (BFL-recommended).
job: extension
config:
  name: "numzoo_klein_lora"          # output folder/filename name
  process:
    - type: 'sd_trainer'
      training_folder: "/root/ai-toolkit/modal_output"   # must match MOUNT_DIR in run_modal.py
      device: cuda:0
      # Captions already start with "NUMZOO." so we do NOT set trigger_word here
      # (avoids double-injecting it). Sample prompts below include NUMZOO explicitly.
      network:
        type: "lora"
        linear: 32
        linear_alpha: 32
      save:
        dtype: float16
        save_every: 250            # checkpoints at 250,500,...,1500 — pick the best
        max_step_saves_to_keep: 6
      datasets:
        - folder_path: "/root/ai-toolkit/numzoo-dataset"
          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: 1500                 # 54 imgs; style LoRA peaks ~750–1500
        gradient_accumulation_steps: 1
        train_unet: true
        train_text_encoder: false
        gradient_checkpointing: true
        noise_scheduler: "flowmatch"
        timestep_type: "weighted"   # klein-recommended (UI default)
        optimizer: "adamw8bit"
        lr: 1e-4
        ema_config:
          use_ema: true
          ema_decay: 0.99
        dtype: bf16
      model:
        name_or_path: "black-forest-labs/FLUX.2-klein-base-4B"
        arch: "flux2_klein_4b"      # resolves to Flux2Klein4BModel (TE=Qwen3-4B, VAE auto)
        quantize: true              # 8-bit to fit comfortably on the GPU
        quantize_te: true
        qtype: "qfloat8"
        low_vram: true
      sample:
        sampler: "flowmatch"
        sample_every: 250
        width: 512                  # match the app's reward image size
        height: 512
        guidance_scale: 4           # base (non-distilled) model uses real CFG
        sample_steps: 25            # base model needs more steps than distilled
        seed: 42
        walk_seed: true
        prompts:
          - "NUMZOO. A cute bunny in an enchanted mushroom forest"
          - "NUMZOO. A cute baby panda on a snowy mountain top"
          - "NUMZOO. A cute puppy and kitten in a cosy cottage garden"
          - "NUMZOO. A cute baby lion, baby tiger and bunny surrounded by sparkling stars"
          - "NUMZOO. A cute unicorn under a rainbow and on a glowing crescent moon"
          - "NUMZOO. A cute baby fox on a sunny beach with ocean waves"
meta:
  name: "[name]"
  version: '1.0'