# Template for tests/local/models.toml (git-ignored: it holds local paths).
#
#   cp tests/local/models.example.toml tests/local/models.toml
#   # edit the paths, delete the [[models]] entries you cannot run
#   python tests/local/regression_te.py record     # with the current dependency versions
#   ... upgrade transformers / diffusers ...
#   python tests/local/regression_te.py compare    # or: pytest tests/local
#
# Every [[models]] entry needs `name` (used for the reference file name) and
# `type` (sd / sdxl / sd3 / flux / lumina / hunyuan_image / anima).
# Optional keys for every entry:
#   enabled = false            skip this entry in `record` / `compare` / pytest runs without deleting it
#                              (it still runs when named explicitly with --models)
#   device  = "cuda" | "cpu"   overrides [common].device
#   dtype   = "float32" | "bfloat16" | "float16"
#                              dtype of the text encoders. Default: float32 for
#                              sd / sdxl, bfloat16 for the others (same as a
#                              typical training run).
#   atol / rtol                tolerances for numpy.allclose on float outputs.
#                              Defaults: fp32 1e-4 / 1e-3, bf16 5e-2 / 3e-2, fp16 2e-2 / 2e-2
#   tokenizer_cache_dir        same meaning as --tokenizer_cache_dir of the trainers
# Token ids and attention masks must always match exactly.
# The diffusers noise schedulers ("schedulers") are always checked; no weights needed.

[common]
device = "cuda"
# reference_dir = "tests/local/references"   # default; git-ignored

# ---- Stable Diffusion 1.x -------------------------------------------------
[[models]]
name = "sd15"
type = "sd"
ckpt = "D:/models/sd/v1-5-pruned-emaonly.safetensors"
v2 = false
clip_skip = 2            # optional; omit for the last layer
max_token_length = 225   # optional: 75 (default) / 150 / 225; >75 exercises chunking
check_vae = true         # encode/decode a fixed image with the diffusers AutoencoderKL

# ---- Stable Diffusion 2.x -------------------------------------------------
[[models]]
name = "sd21"
type = "sd"
ckpt = "D:/models/sd/v2-1_768-ema-pruned.safetensors"
v2 = true
max_token_length = 225

# ---- SDXL -----------------------------------------------------------------
[[models]]
name = "sdxl"
type = "sdxl"
ckpt = "D:/models/sdxl/sd_xl_base_1.0.safetensors"
max_token_length = 225
check_vae = true

# ---- SD3 / SD3.5 ----------------------------------------------------------
[[models]]
name = "sd3"
type = "sd3"
clip_l = "D:/models/sd3/clip_l.safetensors"
clip_g = "D:/models/sd3/clip_g.safetensors"
t5xxl = "D:/models/sd3/t5xxl_fp16.safetensors"
t5xxl_max_token_length = 256
apply_lg_attn_mask = false
apply_t5_attn_mask = false
dtype = "bfloat16"

# ---- FLUX.1 ---------------------------------------------------------------
[[models]]
name = "flux"
type = "flux"
clip_l = "D:/models/flux/clip_l.safetensors"
t5xxl = "D:/models/flux/t5xxl_fp16.safetensors"
t5xxl_max_token_length = 512
apply_t5_attn_mask = false
dtype = "bfloat16"

# ---- Lumina Image 2.0 (Gemma 2 2B) -----------------------------------------
[[models]]
name = "lumina"
type = "lumina"
gemma2 = "D:/models/lumina/gemma_2_2b_fp16.safetensors"
system_prompt = ""       # same as --system_prompt
max_token_length = 256
dtype = "bfloat16"

# ---- HunyuanImage-2.1 (Qwen2.5-VL 7B + byT5) --------------------------------
[[models]]
name = "hunyuan_image"
type = "hunyuan_image"
qwen2_5_vl = "D:/models/hunyuan_image/qwen_2.5_vl_7b.safetensors"
byt5 = "D:/models/hunyuan_image/byt5_small_glyphxl_fp16.safetensors"
dtype = "bfloat16"

# ---- Anima (Qwen3 0.6B + T5 tokenizer) --------------------------------------
[[models]]
name = "anima"
type = "anima"
qwen3 = "D:/models/anima/qwen_3_06b_base.safetensors"   # file (uses configs/qwen3_06b) or a model directory
# t5_tokenizer_path = "D:/models/anima/t5_tokenizer"      # optional; bundled config is used when omitted
qwen3_max_token_length = 512
t5_max_token_length = 512
dtype = "bfloat16"