from .node import BatchedMaskedReferenceGen,ReferenceConditionCombine,MasksToPatches,ConditioningMatchMask,ExtractAttn,PreviewSubjectMask,GenConsistent,MaskToPatchMask,GetCrossSim
NODE_CLASS_MAPPINGS = {

    "BatchedMaskedReferenceGen": BatchedMaskedReferenceGen,
    "MasksToPatches": MasksToPatches,
    "ReferenceConditionCombine": ReferenceConditionCombine,
    "ConditioningMatchMask": ConditioningMatchMask,
    "ExtractAttn": ExtractAttn,
    "PreviewSubjectMask" : PreviewSubjectMask,
    "GenConsistent": GenConsistent,
    "MaskToPatchMask": MaskToPatchMask,
    "GetCrossSim": GetCrossSim,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "ReferenceConditionCombine": "Reference Condition Combine",
    "MasksToPatches": "Masks to Patches",
    "ConditioningMatchMask": "Conditioning Match Mask",
    "BatchedMaskedReferenceGen": "Batched Masked Reference Gen",
    "ExtractAttn": "Extract attention data",
    "PreviewSubjectMask": "Preview Subject Mask",
    "GenConsistent": "GenConsistent",
    "MaskToPatchMask": "MaskToPatchMask",
    "GetCrossSim": "GetCrossSim",
}

__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']