Bjornulf_custom_nodes/clear_vram.py.txt

import torch
import gc

class ClearVRAM:
    def __init__(self):
        self.cleared = False
    
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "clear": ("BOOLEAN", {"default": False, "label": "Clear VRAM"})
            }
        }
    
    RETURN_TYPES = ("STRING",)
    FUNCTION = "clear_vram"
    CATEGORY = "memory_management"
    OUTPUT_NODE = True

    def clear_vram(self, clear):
        if clear and not self.cleared:
            torch.cuda.empty_cache()
            gc.collect()
            torch.cuda.synchronize()
            print("VRAM cleared")
            self.cleared = True
            return ("VRAM cleared",)
        elif not clear:
            self.cleared = False
            return ("Ready to clear VRAM",)
        else:
            return ("VRAM already cleared",)