fluxdev/stable-diffusion-webui-forge

1import time2import torch3import contextlib4from ldm_patched.modules import model_management5from ldm_patched.modules.ops import use_patched_ops6 7 8@contextlib.contextmanager9def automatic_memory_management():10    model_management.free_memory(11        memory_required=3 * 1024 * 1024 * 1024,12        device=model_management.get_torch_device()13    )14 15    module_list = []16 17    original_init = torch.nn.Module.__init__18    original_to = torch.nn.Module.to19 20    def patched_init(self, *args, **kwargs):21        module_list.append(self)22        return original_init(self, *args, **kwargs)23 24    def patched_to(self, *args, **kwargs):25        module_list.append(self)26        return original_to(self, *args, **kwargs)27 28    try:29        torch.nn.Module.__init__ = patched_init30        torch.nn.Module.to = patched_to31        yield32    finally:33        torch.nn.Module.__init__ = original_init34        torch.nn.Module.to = original_to35 36    start = time.perf_counter()37    module_list = set(module_list)38 39    for module in module_list:40        module.cpu()41 42    model_management.soft_empty_cache()43    end = time.perf_counter()44 45    print(f'Automatic Memory Management: {len(module_list)} Modules in {(end - start):.2f} seconds.')46    return47