import comfy.model_management as model_management
import gc
import torch
import time

# Note: This doesn't work with reroute for some reason?
class AnyType(str):
    def __ne__(self, __value: object) -> bool:
        return False

any = AnyType("*")

class UnloadModelNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {"value": (any, )}, # For passthrough
            "optional": {"model": (any, )},
        }
    
    @classmethod
    def VALIDATE_INPUTS(s, **kwargs):
        return True
    
    RETURN_TYPES = (any, )
    FUNCTION = "route"
    CATEGORY = "Unload Model"
    
    def route(self, **kwargs):
        print("Unload Model:")
        loaded_models = model_management.loaded_models()
        if kwargs.get("model") in loaded_models:
            print(" - Model found in memory, unloading...")
            loaded_models.remove(kwargs.get("model"))
        else:
            # Just delete it, I need the VRAM!
            model = kwargs.get("model")
            if type(model) == dict:
                keys = [(key, type(value).__name__) for key, value in model.items()]
                for key, model_type in keys:
                    if key == 'model':
                        print(f"Unloading model of type {model_type}")
                        del model[key]
                        # Emptying the cache after this should free the memory.
        model_management.free_memory(1e30, model_management.get_torch_device(), loaded_models)
        model_management.soft_empty_cache(True)
        try:
            print(" - Clearing Cache...")
            gc.collect()
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
        except:
            print("   - Unable to clear cache")
        return (kwargs.get("value"),)

class UnloadAllModelsNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {"value": (any, )},
        }
    
    @classmethod
    def VALIDATE_INPUTS(s, **kwargs):
        return True
    
    RETURN_TYPES = (any, )
    FUNCTION = "route"
    CATEGORY = "Unload Model"
    
    def route(self, **kwargs):
        print("Unload Model:")
        print(" - Unloading all models...")
        model_management.unload_all_models()
        model_management.soft_empty_cache(True)
        try:
            print(" - Clearing Cache...")
            gc.collect()
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
        except:
            print("   - Unable to clear cache")
        return (kwargs.get("value"),)


NODE_CLASS_MAPPINGS = {
    "UnloadModel": UnloadModelNode,
    "UnloadAllModels": UnloadAllModelsNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "UnloadModel": "Unload Model",
    "UnloadAllModels": "Unload All Models",
}