Function offload_models

src/diffusers/training_utils.py:422–450

Context manager that, if offload=True, moves each module to `device` on enter, then moves it back to its original device on exit. Args: device (`str` or `torch.Device`): Device to move the `modules` to. offload (`bool`): Flag to enable offloading.

(*modules: torch.nn.Module | DiffusionPipeline, device: str | torch.device, offload: bool = True)

@contextmanager
def offload_models(*modules: torch.nn.Module | DiffusionPipeline, device: str | torch.device, offload: bool = True):
    """
    Context manager that, if offload=True, moves each module to `device` on enter, then moves it back to its original
    device on exit.

    Args:
        device (`str` or `torch.Device`): Device to move the `modules` to.
        offload (`bool`): Flag to enable offloading.
    """
    if offload:
        is_model = not any(isinstance(m, DiffusionPipeline) for m in modules)
        # record where each module was
        if is_model:
            original_devices = [next(m.parameters()).device for m in modules]
        else:
            assert len(modules) == 1
            # For DiffusionPipeline, wrap the device in a list to make it iterable
            original_devices = [modules[0].device]
        # move to target device
        for m in modules:
            m.to(device)

    try:
        yield
    finally:
        if offload:
            # move back to original devices
            for m, orig_dev in zip(modules, original_devices):
                m.to(orig_dev)


def parse_buckets_string(buckets_str):