tc5/config.py

import torch

# ─── 1) CONFIG ─────────────────────────────────────────────────────
SAMPLE_RATE = 22050
N_MELS = 80
HOP_LENGTH = 256  # ~86 fps
TIME_SUB = 1
CNN_CH = 128
N_HEADS = 4
D_MODEL = 256
FF_DIM = 512
N_LAYERS = 4
DEPTHWISE_CONV_KERNEL_SIZE = 31
DROPOUT = 0.1
HIDDEN_DIM = 64
N_TYPES = 7
BATCH_SIZE = 4
GRAD_ACCUM_STEPS = 4
LR = 3e-4
EPOCHS = 30
DEVICE = (
    "cuda"
    if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available() else "cpu"
)