(models, input_dir)
| 190 | weights.pop() |
| 191 | |
| 192 | def load_model_hook(models, input_dir): |
| 193 | if cfgs.use_ema: |
| 194 | load_model = EMAModel.from_pretrained(os.path.join(input_dir, "unet_ema"), UNet2DConditionModel) |
| 195 | ema_unet.load_state_dict(load_model.state_dict()) |
| 196 | ema_unet.to(accelerator.device) |
| 197 | del load_model |
| 198 | |
| 199 | for i in range(len(models)): |
| 200 | # pop models so that they are not loaded again |
| 201 | model = models.pop() |
| 202 | |
| 203 | # load diffusers style into model |
| 204 | load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet") |
| 205 | model.register_to_config(**load_model.config) |
| 206 | |
| 207 | model.load_state_dict(load_model.state_dict()) |
| 208 | del load_model |
| 209 | |
| 210 | accelerator.register_save_state_pre_hook(save_model_hook) |
| 211 | accelerator.register_load_state_pre_hook(load_model_hook) |
nothing calls this directly
no test coverage detected