AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00
Code Activity

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

Browse Source 
* 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>
This commit is contained in:
dg845
2024-04-11 10:52:12 -07:00
committed by GitHub
parent 42f25d601a
commit 0bee4d336b
2 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/consistency_distillation/train_lcm_distill_sd_wds.py
View File

@@ -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)

2
examples/consistency_distillation/train_lcm_distill_sdxl_wds.py
View File

@@ -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)