Update config.py for distill

Tencent · Jun 6, 2024 · 3f5bb49 · 3f5bb49
1 parent cb70930
commit 3f5bb49
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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).")