From 3f5bb4924b8389ce4ace71a7f55a66147bdbff39 Mon Sep 17 00:00:00 2001
From: Zhimin Li <46835311+zml-ai@users.noreply.github.com>
Date: Thu, 6 Jun 2024 14:09:34 +0800
Subject: [PATCH] Update config.py for distill

---
 hydit/config.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/hydit/config.py b/hydit/config.py
index 2a731c0..f09a59b 100644
--- a/hydit/config.py
+++ b/hydit/config.py
@@ -19,7 +19,7 @@ def get_args(default_args=None):
     # HunYuan-DiT
     parser.add_argument("--model", type=str, choices=list(HUNYUAN_DIT_CONFIG.keys()), default='DiT-g/2')
     parser.add_argument("--norm", type=str, default="layer", help="Normalization layer type")
-    parser.add_argument("--load-key", type=str, choices=["ema", "module"], default="ema", help="Load model key for HunYuanDiT checkpoint.")
+    parser.add_argument("--load-key", type=str, choices=["ema", "module", "distill"], default="ema", help="Load model key for HunYuanDiT checkpoint.")
     parser.add_argument('--size-cond', type=int, nargs='+', default=[1024, 1024],
                         help="Size condition used in sampling. 2 values are required for height and width. "
                              "If a single value is provided, the image will be treated to (value, value).")