Function no_init_weights

src/diffusers/models/modeling_utils.py:213–229

Context manager to globally disable weight initialization to speed up loading large models. To do that, all the torch.nn.init function are all replaced with skip.

()

@contextmanager
def no_init_weights():
    """
    Context manager to globally disable weight initialization to speed up loading large models. To do that, all the
    torch.nn.init function are all replaced with skip.
    """

    def _skip_init(*args, **kwargs):
        pass

    for name, init_func in TORCH_INIT_FUNCTIONS.items():
        setattr(torch.nn.init, name, _skip_init)
    try:
        yield
    finally:
        # Restore the original initialization functions
        for name, init_func in TORCH_INIT_FUNCTIONS.items():
            setattr(torch.nn.init, name, init_func)


class ModelMixin(torch.nn.Module, PushToHubMixin):