src/model/encoder/vae/config.py

from dataclasses import dataclass
from typing import Tuple


@dataclass
class Config:
    max_stride: int = 64
    lmb_range: Tuple[float, float] = (16.0, 1024.0)
    lmb_limit: float = 8192.0
    lmb_embed_dim: Tuple[int, ...] = (256, 256)
    sin_period: int = 64
    enc_channels: Tuple[int, ...] = (192, 384, 512, 512, 512)
    dec_channels: Tuple[int, ...] = (512, 512, 384, 256, 128)
    # Latent channels for stride-64/32/16/8 blocks.
    # Original default: (32, 32, 96, 8).
    z_dec_channels: Tuple[int, ...] = (32, 32, 64, 16)

    # Test phase
    lmb: float = 1024.0