| { |
| "models": { |
| "encoder": { |
| "name": "SparseUnetVaeEncoder", |
| "args": { |
| "in_channels": 20, |
| "model_channels": [64, 128, 256, 512, 1024], |
| "latent_channels": 32, |
| "num_blocks": [0, 4, 8, 16, 4], |
| "block_type": [ |
| "SparseConvNeXtBlock3d", |
| "SparseConvNeXtBlock3d", |
| "SparseConvNeXtBlock3d", |
| "SparseConvNeXtBlock3d", |
| "SparseConvNeXtBlock3d" |
| ], |
| "down_block_type": [ |
| "SparseResBlockS2C3d", |
| "SparseResBlockS2C3d", |
| "SparseResBlockS2C3d", |
| "SparseResBlockS2C3d" |
| ], |
| "block_args": [{}, {}, {}, {}, {}], |
| "use_fp16": true |
| } |
| }, |
| "decoder": { |
| "name": "SparseUnetVaeDecoder", |
| "args": { |
| "out_channels": 20, |
| "model_channels": [1024, 512, 256, 128, 64], |
| "latent_channels": 32, |
| "num_blocks": [4, 16, 8, 4, 0], |
| "block_type": [ |
| "SparseConvNeXtBlock3d", |
| "SparseConvNeXtBlock3d", |
| "SparseConvNeXtBlock3d", |
| "SparseConvNeXtBlock3d", |
| "SparseConvNeXtBlock3d" |
| ], |
| "up_block_type": [ |
| "SparseResBlockC2S3d", |
| "SparseResBlockC2S3d", |
| "SparseResBlockC2S3d", |
| "SparseResBlockC2S3d" |
| ], |
| "block_args": [{}, {}, {}, {}, {}], |
| "use_fp16": true, |
| "pred_subdiv": false |
| } |
| } |
| }, |
| "dataset": { |
| "name": "SegVaeDataset", |
| "args": { |
| "resolution": 512, |
| "num_classes": 20, |
| "max_tokens": 30000, |
| "min_tokens": 1 |
| } |
| }, |
| "trainer": { |
| "name": "SegVaeTrainer", |
| "args": { |
| "max_steps": 50000, |
| "batch_size_per_gpu": 8, |
| "batch_split": 1, |
| "optimizer": { |
| "name": "AdamW", |
| "args": { |
| "lr": 1e-5, |
| "weight_decay": 0.0 |
| } |
| }, |
| "ema_rate": [0.9999], |
| "fp16_mode": "inflat_all", |
| "fp16_scale_growth": 0.001, |
| "grad_clip": { |
| "name": "AdaptiveGradClipper", |
| "args": { |
| "max_norm": 1.0, |
| "clip_percentile": 95 |
| } |
| }, |
| "finetune_ckpt": { |
| "encoder": "/cos_zw1/share_304064442/hunyuan/easonyzhu/model/TRELLIS.2-4B/ckpts/tex_enc_next_dc_f16c32_fp16.safetensors", |
| "decoder": "/cos_zw1/share_304064442/hunyuan/easonyzhu/model/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16.safetensors" |
| }, |
| "lambda_kl": 1e-5, |
| "i_print": 1, |
| "i_log": 10, |
| "i_sample": 2000, |
| "i_save": 500 |
| } |
| } |
| } |
|
|
