48 lines
1.3 KiB
Python
48 lines
1.3 KiB
Python
STARFLOW_CONFIGS = {
|
|
"starflow_3B_t2i_256x256": {
|
|
"arguments": {
|
|
"no_flip": 1,
|
|
"fsdp": 1,
|
|
"fsdp_text_encoder": 1,
|
|
"img_size": 256,
|
|
"txt_size": 128,
|
|
"channel_size": 4,
|
|
"patch_size": 1,
|
|
"channels": 3072,
|
|
"blocks": 6,
|
|
"layers_per_block": [2, 2, 2, 2, 2, 24],
|
|
"noise_std": 0.3,
|
|
"batch_size": 1024,
|
|
"lr": 6.4e-05,
|
|
"min_lr": 1e-6,
|
|
"nvp": 1,
|
|
"rope": 1,
|
|
"adaln": 0,
|
|
"sos": 1,
|
|
"seq_order": "L2R",
|
|
"wds": 1,
|
|
"use_softplus": 1,
|
|
"cond_top_only": 1,
|
|
"use_final_norm": 1,
|
|
"gradient_checkpoint": 0,
|
|
"vae": "stabilityai/sd-vae-ft-ema",
|
|
"text": "google/flan-t5-xl",
|
|
"cfg": 2.5,
|
|
"drop_label": 0.1,
|
|
"sample_freq": 10,
|
|
"soft_clip": 4.0,
|
|
"latent_norm_regularization": 1e-4,
|
|
}
|
|
},
|
|
}
|
|
|
|
|
|
def get_config_names():
|
|
return sorted(STARFLOW_CONFIGS.keys())
|
|
|
|
|
|
def get_config_args(name):
|
|
if name not in STARFLOW_CONFIGS:
|
|
raise KeyError(f"Unknown STARFlow config preset: {name}")
|
|
return STARFLOW_CONFIGS[name]["arguments"]
|