The larger models truncate to 2048 dims and don't improve results over the native 1.7B that PixelGen was trained with. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
39 lines
825 B
Python
39 lines
825 B
Python
"""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"]
|