Function load_model_hook

train_ae.py:413–426

(models, input_dir)

                weights.pop()

    def load_model_hook(models, input_dir):
        for _ in range(len(models)):
            # pop models so that they are not loaded again
            model = models.pop()

            # load diffusers style into model
            if isinstance(model, VAEClass):
                load_model = VAEClass.from_pretrained(input_dir, subfolder="vae")
            else:
                raise ValueError(f"Unknown model type {model}")

            model.register_to_config(**load_model.config)  # noqa
            model.load_state_dict(load_model.state_dict())
            del load_model

    accelerator.register_save_state_pre_hook(save_model_hook)
    accelerator.register_load_state_pre_hook(load_model_hook)