| { | |
| "model_type": "grinpress-rectified-flow-unet", | |
| "total_params": 33860000, | |
| "unet": { | |
| "in_channels": 3, | |
| "out_channels": 3, | |
| "base_channels": 64, | |
| "channel_mults": [1, 2, 3, 4], | |
| "num_res_blocks": 2, | |
| "attn_levels": [2, 3], | |
| "attn_heads": 8, | |
| "mid_block": ["ResBlock", "AttnBlock", "ResBlock"], | |
| "time_embed_dim": 512, | |
| "norm": "GroupNorm(32)", | |
| "context_dim": 256, | |
| "params": 32230000 | |
| }, | |
| "text_encoder": { | |
| "tokenizer": "ascii_bytes", | |
| "vocab_size": 128, | |
| "pad_id": 0, | |
| "max_len": 64, | |
| "dim": 256, | |
| "layers": 2, | |
| "heads": 4, | |
| "ffn_mult": 4, | |
| "positional": "sinusoidal_1d", | |
| "cfg_dropout": 0.1, | |
| "null": "learned_sequence", | |
| "params": 1630000 | |
| }, | |
| "flow": { | |
| "type": "rectified_flow", | |
| "interpolation": "x_t = (1 - t) * x0 + t * noise", | |
| "target": "v = noise - x0", | |
| "t_sampling": "logit_normal", | |
| "time_input_scale": 1000 | |
| }, | |
| "training": { | |
| "resolution": 128, | |
| "dataset_size": 50000, | |
| "steps": 60000, | |
| "batch_size": 64, | |
| "optimizer": "AdamW", | |
| "lr": 0.0002, | |
| "betas": [0.9, 0.99], | |
| "weight_decay": 0.01, | |
| "lr_schedule": "linear warmup 1000 steps + cosine to 5% of lr", | |
| "grad_clip": 1.0, | |
| "ema_decay": 0.9999, | |
| "precision": "bf16" | |
| }, | |
| "sampling": { | |
| "solver": "euler", | |
| "default_steps": 50, | |
| "default_cfg_scale": 5.0 | |
| }, | |
| "weights_file": "ema_final_bf16.pt", | |
| "image_range": [-1, 1] | |
| } | |