"""
Configuration for FlowMatchRelay — HuggingFace compatible.

Usage:
    from transformers import AutoConfig
    config = AutoConfig.from_pretrained("AbstractPhil/geolip-diffusion-proto", trust_remote_code=True)
"""

from transformers import PretrainedConfig


class FlowMatchRelayConfig(PretrainedConfig):
    model_type = "flow_match_relay"

    def __init__(
        self,
        in_channels=3,
        base_channels=64,
        channel_mults=(1, 2, 4),
        n_classes=10,
        cond_dim=256,
        use_relay=True,
        relay_patch_dim=16,
        relay_n_anchors=16,
        relay_n_phases=3,
        relay_pw_hidden=32,
        relay_gate_init=-3.0,
        relay_mode="channel",
        image_size=32,
        n_sample_steps=50,
        **kwargs,
    ):
        self.in_channels = in_channels
        self.base_channels = base_channels
        self.channel_mults = list(channel_mults)
        self.n_classes = n_classes
        self.cond_dim = cond_dim
        self.use_relay = use_relay
        self.relay_patch_dim = relay_patch_dim
        self.relay_n_anchors = relay_n_anchors
        self.relay_n_phases = relay_n_phases
        self.relay_pw_hidden = relay_pw_hidden
        self.relay_gate_init = relay_gate_init
        self.relay_mode = relay_mode
        self.image_size = image_size
        self.n_sample_steps = n_sample_steps
        super().__init__(**kwargs)