Unverified Commit 5dbe4f5d authored by Junsong Chen's avatar Junsong Chen Committed by GitHub
Browse files

[fix SANA-Sprint] (#11142)



* fix bug in sana conversion script;

* add more model paths;

---------
Co-authored-by: default avatarSayak Paul <spsayakpaul@gmail.com>
parent 1d37f420
...@@ -27,7 +27,10 @@ from diffusers.utils.import_utils import is_accelerate_available ...@@ -27,7 +27,10 @@ from diffusers.utils.import_utils import is_accelerate_available
CTX = init_empty_weights if is_accelerate_available else nullcontext CTX = init_empty_weights if is_accelerate_available else nullcontext
ckpt_ids = [ ckpt_ids = [
"Efficient-Large-Model/Sana_Sprint_0.6B_1024px/checkpoints/Sana_Sprint_0.6B_1024px.pth"
"Efficient-Large-Model/Sana_Sprint_1.6B_1024px/checkpoints/Sana_Sprint_1.6B_1024px.pth"
"Efficient-Large-Model/SANA1.5_4.8B_1024px/checkpoints/SANA1.5_4.8B_1024px.pth", "Efficient-Large-Model/SANA1.5_4.8B_1024px/checkpoints/SANA1.5_4.8B_1024px.pth",
"Efficient-Large-Model/SANA1.5_1.6B_1024px/checkpoints/SANA1.5_1.6B_1024px.pth",
"Efficient-Large-Model/Sana_1600M_4Kpx_BF16/checkpoints/Sana_1600M_4Kpx_BF16.pth", "Efficient-Large-Model/Sana_1600M_4Kpx_BF16/checkpoints/Sana_1600M_4Kpx_BF16.pth",
"Efficient-Large-Model/Sana_1600M_2Kpx_BF16/checkpoints/Sana_1600M_2Kpx_BF16.pth", "Efficient-Large-Model/Sana_1600M_2Kpx_BF16/checkpoints/Sana_1600M_2Kpx_BF16.pth",
"Efficient-Large-Model/Sana_1600M_1024px_MultiLing/checkpoints/Sana_1600M_1024px_MultiLing.pth", "Efficient-Large-Model/Sana_1600M_1024px_MultiLing/checkpoints/Sana_1600M_1024px_MultiLing.pth",
...@@ -314,7 +317,6 @@ def main(args): ...@@ -314,7 +317,6 @@ def main(args):
# SCM Scheduler for Sana Sprint # SCM Scheduler for Sana Sprint
scheduler_config = { scheduler_config = {
"num_train_timesteps": 1000,
"prediction_type": "trigflow", "prediction_type": "trigflow",
"sigma_data": 0.5, "sigma_data": 0.5,
} }
...@@ -378,7 +380,8 @@ if __name__ == "__main__": ...@@ -378,7 +380,8 @@ if __name__ == "__main__":
choices=[ choices=[
"SanaMS_1600M_P1_D20", "SanaMS_1600M_P1_D20",
"SanaMS_600M_P1_D28", "SanaMS_600M_P1_D28",
"SanaMS_4800M_P1_D60", "SanaMS1.5_1600M_P1_D20",
"SanaMS1.5_4800M_P1_D60",
"SanaSprint_1600M_P1_D20", "SanaSprint_1600M_P1_D20",
"SanaSprint_600M_P1_D28", "SanaSprint_600M_P1_D28",
], ],
...@@ -421,7 +424,7 @@ if __name__ == "__main__": ...@@ -421,7 +424,7 @@ if __name__ == "__main__":
"cross_attention_dim": 2240, "cross_attention_dim": 2240,
"num_layers": 20, "num_layers": 20,
}, },
"SanaMS1.5__4800M_P1_D60": { "SanaMS1.5_4800M_P1_D60": {
"num_attention_heads": 70, "num_attention_heads": 70,
"attention_head_dim": 32, "attention_head_dim": 32,
"num_cross_attention_heads": 20, "num_cross_attention_heads": 20,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment