| 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 | |