stage1_layout_flow.json

{
  "name": "TriSparseStructureFlowModel",
  "args": {
    "resolution": 16,
    "channels_per_stream": 8,
    "num_streams": 2,
    "model_channels": 1536,
    "cond_channels": 1024,
    "num_blocks": 30,
    "num_heads": 12,
    "mlp_ratio": 5.3334,
    "pe_mode": "rope",
    "share_mod": true,
    "qk_rms_norm": true,
    "qk_rms_norm_cross": true,
    "dtype": "bfloat16",
    "weight_share": "separate",
    "cross_volume_at": [
      6,
      12,
      18,
      24,
      29
    ],
    "per_stream_io": true,
    "rope_stream_offset": 16.0
  }
}