"""dheeren.config — model hyperparameters (must match training)."""
from dataclasses import dataclass, field
from typing import List, Tuple

@dataclass
class _Config:
    # VAE
    vae_image_size: int = 512
    vae_in_channels: int = 3
    vae_latent_dim: int = 4
    vae_hidden_dims: List[int] = field(default_factory=lambda: [128, 256, 512])

    # CLIP
    clip_vocab_size: int = 49408
    clip_embed_dim: int = 512
    clip_num_layers: int = 12
    clip_num_heads: int = 8
    clip_mlp_ratio: int = 4
    clip_max_seq_length: int = 77
    clip_dropout: float = 0.1

    # UNet
    unet_image_size: int = 64
    unet_in_channels: int = 4
    unet_out_channels: int = 4
    unet_model_channels: int = 192
    unet_num_res_blocks: int = 2
    unet_attention_resolutions: Tuple[int, ...] = (16, 8)
    unet_channel_mult: Tuple[int, ...] = (1, 2, 2, 4)
    unet_dropout: float = 0.1
    unet_num_heads: int = 3
    unet_context_dim: int = 512

    # Flow Matching
    num_diffusion_steps: int = 1000

CFG = _Config()