| # 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' | |