LCM Distill Scripts Fix Bug when Initializing Target U-Net (#6848)

* Initialize target_unet from unet rather than teacher_unet so that we correctly add time_embedding.cond_proj if necessary. * Use UNet2DConditionModel.from_config to initialize target_unet from unet's config. --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
huggingface · Apr 11, 2024 · 0bee4d3 · 0bee4d3
1 parent 42f25d6
commit 0bee4d3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/examples/consistency_distillation/train_lcm_distill_sd_wds.py b/examples/consistency_distillation/train_lcm_distill_sd_wds.py
@@ -945,7 +945,7 @@ def main(args):
 
     # 8. Create target student U-Net. This will be updated via EMA updates (polyak averaging).
     # Initialize from (online) unet
-    target_unet = UNet2DConditionModel(**teacher_unet.config)
+    target_unet = UNet2DConditionModel.from_config(unet.config)
     target_unet.load_state_dict(unet.state_dict())
     target_unet.train()
     target_unet.requires_grad_(False)

diff --git a/examples/consistency_distillation/train_lcm_distill_sdxl_wds.py b/examples/consistency_distillation/train_lcm_distill_sdxl_wds.py
@@ -1004,7 +1004,7 @@ def main(args):
 
     # 8. Create target student U-Net. This will be updated via EMA updates (polyak averaging).
     # Initialize from (online) unet
-    target_unet = UNet2DConditionModel(**teacher_unet.config)
+    target_unet = UNet2DConditionModel.from_config(unet.config)
     target_unet.load_state_dict(unet.state_dict())
     target_unet.train()
     target_unet.requires_grad_(False)