(models, input_dir)
| 545 | weights.pop() |
| 546 | |
| 547 | def load_model_hook(models, input_dir): |
| 548 | for _ in range(len(models)): |
| 549 | # pop models so that they are not loaded again |
| 550 | model = models.pop() |
| 551 | |
| 552 | # load diffusers style into model |
| 553 | if isinstance(model, JoDiffusionModel): |
| 554 | load_model = JoDiffusionModel.from_pretrained(input_dir, subfolder="unet") |
| 555 | else: |
| 556 | raise ValueError(f"Unknown model type {model}") |
| 557 | |
| 558 | model.register_to_config(**load_model.config) # noqa |
| 559 | model.load_state_dict(load_model.state_dict()) |
| 560 | del load_model |
| 561 | |
| 562 | accelerator.register_save_state_pre_hook(save_model_hook) |
| 563 | accelerator.register_load_state_pre_hook(load_model_hook) |
nothing calls this directly
no outgoing calls
no test coverage detected