training_config.json

{
  "learning_rate": 0.0001,
  "weight_decay": 0.01,
  "betas": [
    0.9,
    0.95
  ],
  "grad_clip": 1.0,
  "optimizer": "adamw",
  "warmup_steps": 500,
  "total_steps": 417059,
  "max_epochs": 100,
  "lr_scheduler": "constant",
  "min_lr": 1e-06,
  "warmup_type": "linear",
  "gradient_accumulation": 8,
  "gradient_checkpointing": false,
  "compile_mode": null,
  "ema_decay": 0.9999,
  "shift": 3.0,
  "use_snr_weighting": true,
  "snr_gamma": 5.0,
  "use_huber_loss": true,
  "huber_delta": 0.1,
  "logit_normal_sampling": true,
  "logit_mean": 0.0,
  "logit_std": 1.0,
  "enable_lune": true,
  "lune_weight": 0.1,
  "lune_warmup_steps": 1000,
  "lune_dropout": 0.1,
  "lune_mode": "cosine",
  "enable_sol": true,
  "sol_weight": 0.05,
  "sol_warmup_steps": 2000,
  "sol_dropout": 0.1,
  "use_spatial_weighting": false,
  "text_dropout": 0.1,
  "checkpoint_dir": "./checkpoints",
  "save_every_steps": 1562,
  "keep_last_n_steps": 5,
  "save_every_epochs": 1,
  "keep_last_n_epochs": 3,
  "log_every": 20,
  "tensorboard_dir": "./logs",
  "sample_every": 500,
  "sample_prompts": [
    "subject, object, a red cube on a blue sphere",
    "subject, a golden retriever sitting on grass",
    "subject, object, a coffee cup next to a laptop",
    "subject, a woman with short hair smiling"
  ],
  "sample_dir": "./samples",
  "dtype": "torch.bfloat16",
  "hf_repo_id": "AbstractPhil/tiny-flux-deep",
  "upload_every_steps": 1562,
  "upload_every_epochs": 0
}