"""Model configurations and registry for PixelGen."""

DENOISER_CONFIG_XXL = {
    "patch_size": 16,
    "input_size": 512,
    "in_channels": 3,
    "hidden_size": 1536,
    "num_blocks": 16,
    "num_groups": 24,
    "txt_embed_dim": 2048,
    "txt_max_length": 128,
    "bottleneck_dim": 256,
    "num_text_blocks": 4,
}

MODEL_CONFIGS = {
    "PixelGen-XXL-T2I": {
        "repo_id": "zehongma/PixelGen",
        "filename": "PixelGen_XXL_T2I.ckpt",
        "denoiser_config": DENOISER_CONFIG_XXL,
    },
}

TEXT_ENCODER_CONFIGS = {
    "Qwen/Qwen3-1.7B": {"embed_dim": 2048, "max_length": 128},
}

# Populated at init time with scanned local models
AVAILABLE_MODELS = {}
AVAILABLE_TEXT_ENCODERS = {}


def get_available_devices():
    return ["cuda", "cpu"]


def get_dtype_options():
    return ["bfloat16", "float32"]