internals/util/cache.py

import gc

import torch


def clear_cuda_and_gc():
    clear_cuda()
    clear_gc()


def clear_cuda():
    torch.cuda.empty_cache()


def clear_gc():
    gc.collect()


def auto_clear_cuda_and_gc(controlnet):
    def auto_clear_cuda_and_gc_wrapper(func):
        def wrapper(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            except Exception as e:
                controlnet.cleanup()
                clear_cuda_and_gc()
                raise e

        return wrapper

    return auto_clear_cuda_and_gc_wrapper